Attention-based neural networks with branching blocks
A method for receiving training data for training a neural network to perform a machine learning task and for searching for, using the training data, an optimized neural network architecture for performing the machine learning task is described. Searching for the optimized neural network architecture includes: maintaining population data; maintaining threshold data; and repeatedly performing the following operations: selecting one or more candidate architectures from the population data; generating a new architecture from the one or more selected candidate architectures; for the new architecture: training a neural network having the new architecture until termination criteria for the training are satisfied; and determining a final measure of fitness of the neural network having the new architecture after the training; and adding data defining the new architecture and the final measure of fitness for the neural network having the new architecture to the population data.
1 . A system for performing a machine learning task on a network input to generate a network output, the system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to implement:
an attention neural network configured to perform the machine learning task, the attention neural network comprising a plurality of operation blocks, the plurality of operation blocks comprising a plurality of branching attention blocks, each respective branching attention block comprising a respective first branch that comprises a respective first attention layer that comprises a respective first number of attention heads and a respective second branch that comprises a respective second attention layer that comprises a respective second number of attention heads, each respective branching attention block configured to:
receive a branching attention block input;
provide the branching attention block input to the respective first branch included in the respective branching attention block and to the respective second branch included in the respective branching attention block;
generate, by the respective first attention layer included in the respective first branch included in the respective branching attention block, a first attention layer output by applying an attention mechanism on a first attention layer input using each attention head in the respective first number of attention heads included in the respective first attention layer, wherein the first attention layer input is derived from the branching attention block input;
generate, by the respective second attention layer included in the respective second branch included in the respective branching attention block, a second attention layer output by applying an attention mechanism on a second attention layer input using each attention head in the respective second number of attention heads included in the second attention layer, wherein the second attention layer input is derived from the branching attention block input;
determine a branching attention block output based on determining a combination of the first attention layer output and the second attention layer output; and
provide the branching attention block output as input to a subsequent block in the plurality of operation blocks included in the attention neural network,
wherein, within each respective branching attention block, the respective first number of attention heads of the respective first attention layer included in the respective first branch is different from the respective second number of attention heads of the respective second attention layer included in the respective second branch, and
wherein the respective first number of attention heads of the respective first attention layer included in the respective first branch in each respective branching attention block of the plurality of branching attention blocks is a same as the respective first number of attention heads of the respective first attention layer included in the respective first branch in another respective branching attention block of the plurality of branching attention blocks.
2 . The system of claim 1 , wherein at least one of the plurality of operation blocks included in the attention neural network comprises one or more depth-wise separable convolutional neural network layers.
3 . The system of claim 1 , wherein at least one of the plurality of operation blocks included in the attention neural network comprises one or more gated linear unit (GLU) neural network layers.
4 . The system of claim 1 , wherein at least one of the plurality of operation blocks included in the attention neural network comprises one or more swish activation neural network layers.
5 . The system of claim 1 , wherein the respective first attention layer included in the respective first branch is different than the respective second attention layer included in the respective second branch.
6 . The system of claim 1 , wherein determining the combination comprises computing an addition between the first attention layer output and the second attention layer output.
7 . The system of claim 1 , wherein the respective first attention layer is configured to generate the first attention layer output at least in part by applying a multi-head attention mechanism that uses the respective first number of attention heads to the first attention input layer, wherein the respective second attention layer is configured to generate the second attention layer output at least in part by applying a multi-head attention mechanism that uses the respective second number of attention heads to the second attention layer input, and wherein the first number is greater than the second number.
8 . The system of claim 1 , wherein the network input is an input sequence that includes a respective input at each of multiple positions in an input order, the network output is an output sequence that includes a respective output at each of multiple positions in an output order, or both.
9 . The system of claim 1 , wherein the network input comprises text data, image data, or audio data.
10 . A computer-implemented method comprising, at each respective branching attention block in a plurality of branching attention blocks in an attention neural network:
receiving, at the respective branching attention block, a branching attention block input, wherein the respective branching attention block comprises a respective first branch that comprises a respective first attention layer that comprises a respective first number of attention heads and a respective second branch that comprises a respective second attention layer that comprises a respective second number of attention heads;
providing, by the respective branching attention block, the branching attention block input to the respective first branch included in the respective branching attention block and to the respective second branch included in the respective branching attention block;
generating, by the respective first attention layer included in the respective first branch included in the respective branching attention block, a first attention layer output by applying an attention mechanism on a first attention layer input using each attention head in the respective first number of attention heads included in the respective first attention layer, wherein the first attention layer input is derived from the branching attention block input;
generating, by the respective second attention layer included in the respective second branch included in the respective branching attention block, a second attention layer output by applying an attention mechanism on a second attention layer input using each attention head in the respective second number of attention heads included in the respective second attention layer, wherein the second attention layer input is derived from the branching attention block input;
determining, by the respective branching attention block, a branching attention block output based on determining a combination of the first attention layer output and second attention layer output; and
providing, by the branching attention block, the branching attention block output as input to a subsequent block in a plurality of operation blocks included in the attention neural network,
wherein, within each respective branching attention block, the respective first number of attention heads of the respective first attention layer included in the respective first branch is different from the respective second number of attention heads of the respective second attention layer included in the respective second branch, and
wherein the respective first number of attention heads of the respective first attention layer included in the respective first branch in each respective branching attention block of the plurality of branching attention blocks is a same as the respective first number of attention heads of the respective first attention layer included in the respective first branch in another respective branching attention block of the plurality of branching attention blocks.
11 . The method of claim 10 , wherein at least one of the plurality of operation blocks included in the attention neural network comprises one or more depth-wise separable convolutional neural network layers.
12 . The method of claim 10 , wherein at least one of the plurality of operation blocks included in the attention neural network comprises one or more gated linear unit (GLU) neural network layers.
13 . The method of claim 10 , wherein at least one of the plurality of operation blocks included in the attention neural network comprises one or more swish activation neural network layers.
14 . The method of claim 10 , wherein the respective first attention layer included in the respective first branch is different than the respective second attention layer included in the respective second branch.
15 . The method of claim 10 , wherein determining the combination comprises computing an addition between the first attention layer output and the second attention layer output.
16 . The method of claim 14 , wherein the first attention layer is configured to generate the first attention layer output at least in part by applying a multi-head attention mechanism that uses the first number of attention heads to the first attention layer input, wherein the second attention layer is configured to generate the second attention layer output at least in part by applying a multi-head attention mechanism that uses the second number of attention heads to the second attention layer input, and wherein the first number is greater than the second number.
17 . The method of claim 10 , wherein the network input is an input sequence that includes a respective input at each of multiple positions in an input order, the network output is an output sequence that includes a respective output at each of multiple positions in an output order, or both.
18 . The method of claim 10 , wherein the network input comprises text data, image data, or audio data.
19 . One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to implement:
an attention neural network configured to perform the machine learning task, the attention neural network comprising a plurality of operation blocks, the plurality of operation blocks comprising a plurality of branching attention blocks, each respective branching attention block comprising a respective first branch that comprises a respective first attention layer that comprises a respective first number of attention heads and a respective second branch that comprises a respective second attention layer that comprises a respective second number of attention heads, each respective branching attention block configured to:
receive a branching attention block input;
provide the branching attention block input to the respective first branch included in the respective branching attention block and to the respective second branch included in the respective branching attention block;
generate, by the respective first attention layer included in the respective first branch included in the respective branching attention block, a first attention layer output by applying an attention mechanism on a first attention layer input using each attention head in the respective first number of attention heads included in the respective first attention layer, wherein the first attention layer input is derived from the branching attention block input;
generate, by the respective second attention layer included in the respective second branch included in the respective branching attention block, a second attention layer output by applying an attention mechanism on a second attention layer input using each attention head in the respective second number of attention heads included in the second attention layer, wherein the second attention layer input is derived from the branching attention block input;
determine a branching attention block output based on determining a combination of the first attention layer output and the second attention layer output; and
provide the branching attention block output as input to a subsequent block in the plurality of operation blocks included in the attention neural network,
wherein, within each respective branching attention block, the respective first number of attention heads of the respective first attention layer included in the respective first branch is different from the respective second number of attention heads of the respective second attention layer included in the respective second branch, and
wherein the respective first number of attention heads of the respective first attention layer included in the respective first branch in each respective branching attention block of the plurality of branching attention blocks is a same as the respective first number of attention heads of the respective first attention layer included in the respective first branch in another respective branching attention block of the plurality of branching attention blocks.
20 . The computer storage media of claim 19 , wherein the network input comprises text data, image data, or audio data.