PUMA: margin-based data pruning

TL;DR

PUMA tackles the data-inefficiency of adversarial training under large-scale synthetic data by introducing margin-based data pruning. It uses DeepFool to estimate per-sample model margin, pruning high-margin (less informative) samples while adapting the per-sample attack strength $oldsymbol{}_i$ for low-margin samples, thereby improving the accuracy-robustness trade-off with less data. The approach is supported by a perceptron toy analysis and extensive image-classification experiments (CIFAR10 with EDM augmentation and ImageNet-21K variants), showing substantial accuracy gains over state-of-the-art robustness methods without sacrificing robustness at practical pruning levels. These results underscore data quality (margin) as a critical factor in robust learning and open avenues for margin-guided data generation and efficient adversarial training. Key technical elements include the use of a margin model via DeepFool, adaptive per-sample perturbation norms, and empirical demonstrations across multiple architectures and data scales, highlighting PUMA’s potential to reduce computational costs while maintaining or enhancing robustness and accuracy.

Abstract

Deep learning has been able to outperform humans in terms of classification accuracy in many tasks. However, to achieve robustness to adversarial perturbations, the best methodologies require to perform adversarial training on a much larger training set that has been typically augmented using generative models (e.g., diffusion models). Our main objective in this work, is to reduce these data requirements while achieving the same or better accuracy-robustness trade-offs. We focus on data pruning, where some training samples are removed based on the distance to the model classification boundary (i.e., margin). We find that the existing approaches that prune samples with low margin fails to increase robustness when we add a lot of synthetic data, and explain this situation with a perceptron learning task. Moreover, we find that pruning high margin samples for better accuracy increases the harmful impact of mislabeled perturbed data in adversarial training, hurting both robustness and accuracy. We thus propose PUMA, a new data pruning strategy that computes the margin using DeepFool, and prunes the training samples of highest margin without hurting performance by jointly adjusting the training attack norm on the samples of lowest margin. We show that PUMA can be used on top of the current state-of-the-art methodology in robustness, and it is able to significantly improve the model performance unlike the existing data pruning strategies. Not only PUMA achieves similar robustness with less data, but it also significantly increases the model accuracy, improving the performance trade-off.

Figures (9)

• Figure 1: Correlation error between the teacher and the student weights relative to $\alpha$ when standardly training the student. The orange line shows the performance when not pruning. For $\alpha < 2$, only the PD strategy improves performance. For $\alpha > 2$, as we increase the pruning ratio, the PE strategy outperforms significantly not pruning and has exponential scaling instead of power law.
• Figure 2: Correlation error between the teacher and the student weights relative to $\alpha$ when adversarially training the student with $\varepsilon = 0.001$.
• Figure 3: Correlation error between the teacher and student weights relative to $\alpha$ when adversarially training the student with $\varepsilon = 0.01$.
• Figure 4: Correlation error between the teacher and the student weights relative to $\alpha$ when adversarially training the student with $\varepsilon = 0.01$, after filtering all the samples with margin lower than $\varepsilon$.
• Figure 5: PD strategy for the ResNet-18 model adversarially trained on CIFAR10 with 1M EDM-generated samples. At the beginning (i.e. first 100-150 epochs), the model still has not seen enough new images and pruning improves robustness. But as the model trains with more new images, its robustness decreases compared to random pruning. Moreover, accuracy is significantly worse in all training stages.