Processing images using mixture of experts
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for generating predictions about images. One of the systems includes a neural network comprising a sequence of one or more network blocks that are each configured to perform operations comprising: obtaining a block input that represents an intermediate representation of an input image; determining a plurality of patches of the block input or of an updated representation of the block input, wherein each patch comprises a different subset of elements of the block input or of the updated representation of the block input; assigning each patch to one or more respective expert modules of a plurality of expert modules of the network block; for each patch of the plurality of patches, processing the patch using the corresponding expert modules to generate respective module outputs; and generating a block output by combining the module outputs.
1 . A system comprising a neural network that is configured to process an input image and to generate a network output for the input image, the neural network comprising a sequence of one or more network blocks that are each configured to perform operations comprising:
obtaining a block input that represents an intermediate representation of the input image;
determining a plurality of patches of the block input or of an updated representation of the block input, wherein each patch comprises a different subset of elements of the block input or of the updated representation of the block input;
assigning each patch to one or more respective expert modules of a plurality of expert modules of the network block;
for each patch of the plurality of patches, processing the patch using the one or more corresponding expert modules to generate one or more respective module outputs; and
generating a block output that is an updated representation of the input image and that has a same dimensionality as the block input by combining the module outputs generated by the respective expert modules.
2 . The system of claim 1 , wherein assigning each patch to one or more respective experts module of the network block comprises:
processing the patch using a routing subnetwork to generate, for each expert module of the plurality of expert modules, a respective routing score;
determining one or more highest routing scores of the plurality of routing scores; and
assigning the patch to the expert module corresponding to each of the one or more highest routing scores.
3 . The system of claim 2 , wherein processing the patch using a routing subnetwork to generate, for each expert module of the plurality of expert modules, a respective routing score comprises:
processing the patch using one or more feedforward neural network layers to generate a respective first intermediate routing score for each expert module;
sampling, for each expert module, a noise value and adding the noise value to the first intermediate routing score of the expert module to generate a second intermediate routing score for the expert module;
applying, for each expert module, a non-linear activation function to the second intermediate routing score of the expert module to generate the routing score for the expert module.
4 . The system of claim 2 , wherein, for each patch of the plurality of patches, processing the patch using the one or more corresponding expert modules to generate one or more respective module outputs comprises:
processing, for each of the one or more expert modules corresponding to the patch, the patch using the expert module to generate a respective initial module output; and
scaling, for each of the one or more expert modules corresponding to the patch, the corresponding initial module output according to the routing score corresponding to the expert module and the patch.
5 . The system of claim 2 , wherein the neural network has been trained using one or more of:
an importance loss that encourages the routing scores corresponding to each expert module to have a same measure of central tendency, or
a load loss that encourages respective likelihoods that each expert module is assigned a patch to have a same measure of central tendency.
6 . The system of claim 5 , wherein determining the importance loss for a batch of training images X comprises:
for each expert module i, computing:
imp
i
(
X
)
=
∑
x
∈
X
f
i
(
x
)
wherein f i is a function defined by the routing subnetwork for computing the routing score for the expert module i or an intermediate routing score for the expert module i; and
computing the importance loss to encourage, for each expert module i, imp i (X) to have a same measure of central tendency.
7 . The system of claim 6 , wherein the importance loss is:
ℒ
imp
(
X
)
=
(
std
(
imp
1
:
E
(
X
)
)
mean
(
imp
1
:
E
(
X
)
)
)
2
wherein E is a number of the plurality of expert modules.
8 . The system of claim 5 , wherein:
each patch is assigned to k different expert modules, k≥1, and
determining the load loss for a batch of training images X comprises:
for each training image x in the batch of training images X, computing:
threshold
k
(
x
)
=
max
(
k
)
g
i
(
x
)
wherein g i (x) is a function defined by the routing subnetwork for computing the routing score for the expert module i or an intermediate routing score for the expert module i;
for each training image x and for each expert module i, computing:
p
i
(
x
)
=
P
(
g
i
(
x
)
≥
threshold
k
(
x
)
)
wherein P(e) is a likelihood of event e occurring; and
for each expert module i, computing:
load
i
(
X
)
=
∑
x
∈
X
p
i
(
x
)
and computing the load loss to encourage, for each expert module i, load i (x) to have a same measure of central tendency.
9 . The system of claim 8 , wherein the load loss is:
ℒ
load
(
X
)
=
(
std
(
load
1
:
E
(
X
)
)
mean
(
load
1
:
E
(
X
)
)
)
2
wherein E is a number of the plurality of expert modules.
10 . The system of claim 1 , wherein processing, for each patch of the plurality of patches, the patch using the corresponding expert module comprises processing the patches in parallel across the expert modules.
11 . The system of claim 1 , wherein a number of the plurality of expert modules is greater than a number of the plurality of patches.
12 . The system of claim 1 , wherein the block output of a final network block in the sequence of network blocks is the network output, or wherein the block output of the final network block in the sequence of network blocks is further processed using one or more neural network layers to generate the network output.
13 . The system of claim 1 , wherein:
the network output is a classification output that includes a respective score corresponding to each of multiple categories,
the network output is an element-level classification output that includes, for each of one or more elements in the image, a respective score corresponding to each of multiple categories,
the network output is a regression output that estimates one or more continuous variables that characterize the image, or
the neural network is configured to process multiple images that are video frames to generate an output that characterizes the video frames.
14 . One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one more computers to implement a neural network that is configured to process an input image and to generate a network output for the input image, the neural network comprising a sequence of one or more network blocks that are each configured to perform operations comprising:
obtaining a block input that represents an intermediate representation of the input image;
determining a plurality of patches of the block input or of an updated representation of the block input, wherein each patch comprises a different subset of elements of the block input or of the updated representation of the block input;
assigning each patch to one or more respective expert modules of a plurality of expert modules of the network block;
for each patch of the plurality of patches, processing the patch using the one or more corresponding expert modules to generate one or more respective module outputs; and
generating a block output that is an updated representation of the input image and that has a same dimensionality as the block input by combining the module outputs generated by the respective expert modules.
15 . A method performed by one or more computers, the method comprising:
receiving an input image; and
processing the input image using a neural network to generate a network output for the input image, the neural network comprising a sequence of one or more network blocks that are each configured to perform operations comprising:
obtaining a block input that represents an intermediate representation of the input image;
determining a plurality of patches of the block input or of an updated representation of the block input, wherein each patch comprises a different subset of elements of the block input or of the updated representation of the block input;
assigning each patch to one or more respective expert modules of a plurality of expert modules of the network block;
for each patch of the plurality of patches, processing the patch using the one or more corresponding expert modules to generate one or more respective module outputs; and
generating a block output that is an updated representation of the input image and that has a same dimensionality as the block input by combining the module outputs generated by the respective expert modules.
16 . The method of claim 15 , wherein assigning each patch to one or more respective experts module of the network block comprises:
processing the patch using a routing subnetwork to generate, for each expert module of the plurality of expert modules, a respective routing score;
determining one or more highest routing scores of the plurality of routing scores; and
assigning the patch to the expert module corresponding to each of the one or more highest routing scores.
17 . The method of claim 16 , wherein processing the patch using a routing subnetwork to generate, for each expert module of the plurality of expert modules, a respective routing score comprises:
processing the patch using one or more feedforward neural network layers to generate a respective first intermediate routing score for each expert module;
sampling, for each expert module, a noise value and adding the noise value to the first intermediate routing score of the expert module to generate a second intermediate routing score for the expert module;
applying, for each expert module, a non-linear activation function to the second intermediate routing score of the expert module to generate the routing score for the expert module.
18 . The method of claim 17 , wherein, for each patch of the plurality of patches, processing the patch using the one or more corresponding expert modules to generate one or more respective module outputs comprises:
processing, for each of the one or more expert modules corresponding to the patch, the patch using the expert module to generate a respective initial module output; and
scaling, for each of the one or more expert modules corresponding to the patch, the corresponding initial module output according to the routing score corresponding to the expert module and the patch.
19 . The method of claim 17 , wherein the neural network has been trained using one or more of:
an importance loss that encourages the routing scores corresponding to each expert module to have a same measure of central tendency, or
a load loss that encourages respective likelihoods that each expert module is assigned a patch to have a same measure of central tendency.
20 . The method of claim 19 , wherein determining the importance loss for a batch of training images X comprises:
for each expert module i, computing:
imp
i
(
X
)
=
∑
x
∈
X
f
i
(
x
)
wherein f i is a function defined by the routing subnetwork for computing the routing score for the expert module i or an intermediate routing score for the expert module i; and
computing the importance loss to encourage, for each expert module i, imp i (X) to have a same measure of central tendency.