Learning to select vocabularies for categorical features

A method performed by one or more data processing apparatus for determining, for each of one or more categorical feature values, a respective embedding dimensionality of an embedding that should be generated for the categorical feature value during processing of inputs by a machine learning model having a plurality of machine learning model parameters, the method comprising:

• generating, using a controller neural network having a plurality of controller parameters and in accordance with current values of the controller parameters, a batch of output sequences, each output sequence in the batch specifying, for each categorical feature value, a respective embedding dimensionality of an embedding that should be generated for the categorical feature value during processing of inputs by the machine learning model;
• for each output sequence in the batch, determining a performance metric of the machine learning model on a machine learning task after the machine learning model has been trained to perform the machine learning task when the machine learning model generates embeddings of the categorical feature values with the embedding dimensionalities specified by the output sequence, wherein during the training:
• the machine learning model is configured to process an input that comprises one or more categorical feature values by performing operations comprising:
• mapping each categorical feature value included in the input to an embedding having the corresponding embedding dimensionality specified by the output sequence;
• and
• processing the embeddings to generate a machine learning model output for the input;
• and
• using the performance metrics for the output sequences in the batch to adjust the current values of the controller parameters of the controller neural network.

A system comprising:

• one or more computers;
• and
• one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations for determining, for each of one or more categorical features, a respective vocabulary of categorical feature values of the categorical feature that should be active during processing of inputs by a machine learning model having a plurality of machine learning model parameters, the operations comprising:
• generating, using a controller neural network having a plurality of controller parameters and in accordance with current values of the controller parameters, a batch of output sequences, each output sequence in the batch specifying, for each of the categorical features, a respective vocabulary of categorical feature values of the categorical feature that:
• (i) should be active during processing of inputs by the machine learning model, and (ii) is a subset of a discrete set of possible categorical feature values of the categorical feature;
• for each output sequence in the batch, determining a performance metric of the machine learning model on a machine learning task after the machine learning model has been trained to perform the machine learning task with only the respective vocabulary of categorical feature values of each categorical feature specified by the output sequence being active, wherein during the training:
• the machine learning model is configured to process an input that comprises one or more possible categorical feature values of respective categorical features by performing operations comprising:
• for only those possible categorical feature values included in the input that are specified as active by the output sequence, mapping the possible categorical feature value to a corresponding embedding that is iteratively adjusted during the training;
• and
• processing the embeddings to generate a machine learning model output for the input;
• and
• using the performance metrics for the output sequences in the batch to adjust the current values of the controller parameters of the controller neural network.