Single-stage model training for neural architecture search

One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:

• receiving training data for a particular machine learning task;
• generating trained values for a shared set of parameters by training a plurality of neural networks jointly on the training data, wherein each of the plurality of neural networks has parameters that are a subset of the shared set, wherein each of the plurality of neural networks has a respective architecture selected from a search space of different architectures, and wherein the search space is defined by a respective set of possible values for each of a plurality of architectural dimensions, and wherein the joint training comprises performing the following at each of a plurality of training steps:
• obtaining training data for the training step;
• sampling a plurality of neural networks of different sizes from the plurality of neural networks being trained, wherein each sampled neural network has a different subset of the shared set of parameters;
• training the sampled neural networks on the training data for the training step to determine, for each sampled neural network, gradients with respect the respective subset of the shared set of parameters that are included in the sampled neural network;
• and
• aggregating the gradients determined for the set of neural networks to determine an update to the shared set of parameters, comprising, for each parameter that is shared between two or more of the sampled neural networks, aggregating the gradients with respect to the parameter for the two or more sampled neural networks;
• determining a respective performance benchmark on the machine learning task for each of a subset of neural networks, each neural network in the subset having a different architecture selected from a proper subset of the search space of different architectures;
• receiving data specifying one or more constraints on computational resource consumption for performing inference on one or more inputs for the particular machine learning task on a particular set of one or more target hardware devices;
• and
• selecting a neural network to be deployed for performing the machine learning task, comprising:
• identifying, from the neural networks in the subset that satisfy the one or more constraints, an initial neural network based on the respective performance benchmarks;
• generating a plurality of modified neural networks by perturbing the architecture of the identified initial neural network;
• determining a respective performance benchmark on the particular machine learning task for each of the modified neural networks, and
• selecting, from the identified neural network and the plurality of modified neural networks generated by perturbing the architecture of the identified initial neural network, a neural network based on the respective performance benchmarks;
• and
• providing the selected neural network for performing inference on the particular set of one or more target hardware devices.

A system comprising one or more computers and one or more storage devices storing instructions that when executed by one or more computers are operable to cause the one or more computers to perform operations comprising:

• receiving training data for a particular machine learning task;
• generating trained values for a shared set of parameters by training a plurality of neural networks jointly on the training data, wherein each of the plurality of neural networks has parameters that are a subset of the shared set, wherein each of the plurality of neural networks has a respective architecture selected from a search space of different architectures, and wherein the search space is defined by a respective set of possible values for each of a plurality of architectural dimensions, and wherein the joint training comprises performing the following at each of a plurality of training steps:
• obtaining training data for the training step;
• sampling a plurality of neural networks of different sizes from the plurality of neural networks being trained, wherein each sampled neural network has a different subset of the shared set of parameters;
• training the sampled neural networks on the training data for the training step to determine, for each sampled neural network, gradients with respect the respective subset of the shared set of parameters that are included in the sampled neural network;
• and
• aggregating the gradients determined for the set of neural networks to determine an update to the shared set of parameters, comprising, for each parameter that is shared between two or more of the sampled neural networks, aggregating the gradients with respect to the parameter for the two or more sampled neural networks;
• determining a respective performance benchmark on the machine learning task for each of a subset of neural networks, each neural network in the subset having a different architecture selected from a proper subset of the search space of different architectures;
• receiving data specifying one or more constraints on computational resource consumption for performing inference on one or more inputs for the particular machine learning task on a particular set of one or more target hardware devices;
• and
• selecting a neural network to be deployed for performing the machine learning task, comprising:
• identifying, from the neural networks in the subset that satisfy the one or more constraints, an initial neural network based on the respective performance benchmarks;
• generating a plurality of modified neural networks by perturbing the architecture of the identified initial neural network;
• determining a respective performance benchmark on the particular machine learning task for each of the modified neural networks, and
• selecting, from the identified neural network and the plurality of modified neural networks generated by perturbing the architecture of the identified initial neural network, a neural network based on the respective performance benchmarks;
• and
• providing the selected neural network for performing inference on the particular set of one or more target hardware devices.