Methods, systems, and apparatus, including computer programs encoded on computer storage media, for training a machine learning model on training data. In one aspect, one of the methods include: obtaining a training data set comprising a plurality of training inputs; obtaining data defining an original search space of a plurality of candidate data augmentation policies; generating, from the original search space, a compact search space that has one or more global hyperparameters; and training the machine learning model on the training data using one or more final data augmentation policies generated from the compact search space.