Neural networks with switch layers

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 a machine learning task, the attention neural network comprising a plurality of layers, each layer comprising an attention sub-layer and a feed-forward sub-layer, the attention sub-layer configured to:
• receive an input sequence for the layer comprising a respective layer input at each of one or more positions;
• and
• generate an attended input sequence at least in part by applying an attention mechanism to the input sequence for the layer, the attended input sequence comprising a respective attended layer input at each of the one or more positions,
• wherein, for each of one or more of the plurality of layers, the feed-forward layer of the layer includes a respective switch layer that comprises:
• (i) a respective learned routing function having routing parameters, and
• (ii) a respective plurality of expert neural networks each having a respective set of expert parameters,
• wherein the respective switch layer is configured to:
• for each attended layer input in the attended input sequence generated by the attention sub-layer of the layer:
• apply the respective learned routing function in the switch layer to the attended layer input in accordance with current values of the routing parameters of the respective learned routing function to generate a score distribution that includes a respective routing score for each of the plurality of expert neural networks in the switch layer;
• select, from the plurality of expert neural networks, only one of the expert neural networks using the score distribution;
• process the attended layer input using only the selected expert neural network and in accordance with current values of the expert parameters of the selected expert neural network to generate an initial switch layer output for the attended layer input;
• and
• generate a final switch layer output for the switch layer for the attended layer input from the initial switch layer output generated by the selected expert neural network.

A method performed by one or more computers, the method comprising:

• receiving a network input;
• and
• processing the network input using an attention neural network configured to perform a machine learning task to generate a network output for the network input, the attention neural network comprising a plurality of layers, each layer comprising an attention sub-layer and a feed-forward sub-layer, the attention sub-layer configured to:
• receive an input sequence for the layer comprising a respective layer input at each of one or more positions;
• and
• generate an attended input sequence at least in part by applying an attention mechanism to the input sequence for the layer, the attended input sequence comprising a respective attended layer input at each of the one or more positions,
• wherein, for each of one or more of the plurality of layers, the feed-forward layer of the layer includes a respective switch layer that comprises:
• (i) a respective learned routing function having routing parameters, and
• (ii) a respective plurality of expert neural networks each having a respective set of expert parameters,
• wherein the respective switch layer is configured to:
• for each attended layer input in the attended input sequence generated by the attention sub-layer of the layer:
• apply the respective learned routing function in the switch layer to the attended layer input in accordance with current values of the routing parameters of the respective learned routing function to generate a score distribution that includes a respective routing score for each of the plurality of expert neural networks in the switch layer;
• select, from the plurality of expert neural networks, only one of the expert neural networks using the score distribution;
• process the attended layer input using only the selected expert neural network and in accordance with current values of the expert parameters of the selected expert neural network to generate an initial switch layer output for the attended layer input;
• and
• generate a final switch layer output for the switch layer for the attended layer input from the initial switch layer output generated by the selected expert neural network.

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

• receiving a network input;
• and
• processing the network input using an attention neural network configured to perform a machine learning task to generate a network output for the network input, the attention neural network comprising a plurality of layers, each layer comprising an attention sub-layer and a feed-forward sub-layer, the attention sub-layer configured to:
• receive an input sequence for the layer comprising a respective layer input at each of one or more positions;
• and
• generate an attended input sequence at least in part by applying an attention mechanism to the input sequence for the layer, the attended input sequence comprising a respective attended layer input at each of the one or more positions,
• wherein, for each of one or more of the plurality of layers, the feed-forward layer of the layer includes a respective switch layer that comprises:
• (i) a respective learned routing function having routing parameters, and
• (ii) a respective plurality of expert neural networks each having a respective set of expert parameters,
• wherein the respective switch layer is configured to:
• for each attended layer input in the attended input sequence generated by the attention sub-layer of the layer:
• apply the respective learned routing function in the switch layer to the attended layer input in accordance with current values of the routing parameters of the respective learned routing function to generate a score distribution that includes a respective routing score for each of the plurality of expert neural networks in the switch layer;
• select, from the plurality of expert neural networks, only one of the expert neural networks using the score distribution;
• process the attended layer input using only the selected expert neural network and in accordance with current values of the expert parameters of the selected expert neural network to generate an initial switch layer output for the attended layer input;
• and
• generate a final switch layer output for the switch layer for the attended layer input from the initial switch layer output generated by the selected expert neural network.