Multi-stage computationally efficient neural network inference
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for multi-stage computationally-efficient inference using a first and second neural network.
1 . A method performed by one or more computers, the method comprising:
obtaining a first network input;
processing the first network input using a smaller neural network to generate a first network output for the first network input that comprises a respective first score for each of a first plurality of categories;
determining, from the first network output for the first network input, whether the first network output for the first network input is likely to be inaccurate;
in response to determining that the first network output for the first network input is likely to be inaccurate:
providing the first network input as input to a larger neural network that is configured to process the first network input to generate a second network output for the first network input that comprises a respective second score for each of a second plurality of categories, the second plurality of categories comprising all of the categories in the first plurality of categories, wherein the larger neural network is more computationally expensive than the smaller neural network; and
classifying the first network input based on the second network output for the first network input;
obtaining a second network input;
processing the second network input using the smaller neural network to generate a first network output for the second network input that comprises a respective first score for each of the first plurality of categories;
determining, from the first network output for the second network input, whether the first network output for the second network input is likely to be inaccurate; and
in response to determining that the first network output for the second network input is not likely to be inaccurate:
classifying second first network input based on the first network output for the second network input without providing the first network input as input to the larger neural network.
2 . The method of claim 1 , wherein classifying the first network input based on the second network output comprises:
classifying the first network input as belonging to one or more categories having the highest second scores.
3 . The method of claim 1 , further comprising:
in response to determining that the first network output is not likely to be inaccurate:
classifying the first network input based on the first network output without providing the first network input as input to the larger neural network.
4 . The method of claim 1 , wherein the one or more computers are an edge computing device and wherein the smaller neural network is deployed on the edge device.
5 . The method of claim 4 , wherein the smaller neural network is deployed on application-specific integrated circuit (ASIC) or a field-programmable gate array (FPGA) on the edge computing device.
6 . The method of claim 4 , wherein the larger neural network is deployed on one or more second computers remote from the edge computing device, wherein providing the first network input as input to the larger neural network comprises providing the first network input from the edge computing device to the one or more second computers over a data communication network, and wherein the method further comprises obtaining the second network output over the data communication network.
7 . The method of claim 1 , wherein determining, from the first network output, whether the first network output is likely to be inaccurate comprises:
determining whether a category having a highest first score is in a predetermined proper subset of the first plurality of categories; and
determining that the first network output is likely to be inaccurate when the category having the highest first score is in the predetermined proper subset of the plurality of categories.
8 . The method of claim 7 , wherein the predetermined proper subset of the plurality of categories has been selected based on a frequency of occurrence of each category in training data that was used to train the smaller neural network.
9 . The method of claim 1 , wherein determining, from the first network output, whether the first network output is likely to be inaccurate comprises:
determining whether a difference between (i) a highest first score for any category in the first plurality of categories and (ii) a second highest first score for any category in the first plurality of categories satisfies a threshold; and
determining that the first network output is likely to be inaccurate when the difference satisfies the threshold.
10 . The method of claim 1 , wherein the first network output further comprises a first score for an abstain class, and wherein determining, from the first network output, whether the first network output is likely to be inaccurate comprises:
determining whether the first score for the abstain class is higher than the first score for any of the first plurality of categories; and
determining that the first network output is likely to be inaccurate when the first score for the abstain class is higher than the first score for any of the first plurality of categories.
11 . The method of claim 1 , wherein the smaller neural network has been trained using pseudo labels generated by the larger neural network by processing training network inputs.
12 . The method of claim 11 , wherein the smaller neural network has been trained to minimize a loss function that includes a first term that encourages, for any given training network input that has a ground truth label that assigns the given training network input to any of a predetermined proper subset of the second plurality of categories, a first network output generated by the smaller neural network for the given training input to match a pseudo label generated from a second network output generated by the larger neural network for the given training input.
13 . The method of claim 12 , wherein the first term encourages, for any given training network input that has a ground truth label that assigns the given training network input to any category that is not in the predetermined proper subset of the plurality of categories, a first network output generated by the smaller neural network for the given training input to match a distribution that is independent of any second network output generated by the larger neural network for the given training input.
14 . The method of claim 13 , wherein the distribution is a label-smoothed distribution generated from the ground truth label for the given training network input.
15 . The method of claim 13 , wherein the distribution is a uniform distribution over the second plurality of categories.
16 . The method of claim 13 , wherein the distribution is a one hot distribution that assigns a non-zero score to only an abstain class.
17 . The method of claim 12 , wherein the predetermined proper subset of the second plurality of categories has been selected based on a frequency of occurrence of each category in training data that was used to train the smaller neural network.
18 . The method of claim 11 , wherein the smaller neural network has been trained to minimize a loss function that includes a first term that encourages, for any given training network input of a predetermined proper subset of the training inputs in the training data, a first network output generated by the smaller neural network for the given training input to match a pseudo label generated from a second network output generated by the larger neural network for the given training input.
19 . The method of claim 18 , wherein the predetermined proper subset of the training inputs in the training data has been selected based on margins between scores in second network outputs generated for the training inputs in the training data by the larger neural network.
20 . The method of claim 1 , wherein the smaller neural network has fewer parameters than the larger neural network.
21 . A 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 perform operations comprising:
obtaining a first network input;
processing the first network input using a smaller neural network to generate a first network output for the first network input that comprises a respective first score for each of a first plurality of categories;
determining, from the first network output for the first network input, whether the first network output for the first network input is likely to be inaccurate;
in response to determining that the first network output for the first network input is likely to be inaccurate:
providing the first network input as input to a larger neural network that is configured to process the first network input to generate a second network output for the first network input that comprises a respective second score for each of a second plurality of categories, the second plurality of categories comprising all of the categories in the first plurality of categories, wherein the larger neural network is more computationally expensive than the smaller neural network; and
classifying the first network input based on the second network output for the first network input;
obtaining a second network input;
processing the second network input using the smaller neural network to generate a first network output for the second network input that comprises a respective first score for each of the first plurality of categories;
determining, from the first network output for the second network input, whether the first network output for the second network input is likely to be inaccurate; and
in response to determining that the first network output for the second network input is not likely to be inaccurate:
classifying second first network input based on the first network output for the second network input without providing the first network input as input to the larger neural network.
22 . 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 perform operations comprising:
obtaining a first network input;
processing the first network input using a smaller neural network to generate a first network output for the first network input that comprises a respective first score for each of a first plurality of categories;
determining, from the first network output for the first network input, whether the first network output for the first network input is likely to be inaccurate;
in response to determining that the first network output for the first network input is likely to be inaccurate:
providing the first network input as input to a larger neural network that is configured to process the first network input to generate a second network output for the first network input that comprises a respective second score for each of a second plurality of categories, the second plurality of categories comprising all of the categories in the first plurality of categories, wherein the larger neural network is more computationally expensive than the smaller neural network; and
classifying the first network input based on the second network output for the first network input;
obtaining a second network input;
processing the second network input using the smaller neural network to generate a first network output for the second network input that comprises a respective first score for each of the first plurality of categories;
determining, from the first network output for the second network input, whether the first network output for the second network input is likely to be inaccurate; and
in response to determining that the first network output for the second network input is not likely to be inaccurate:
classifying second first network input based on the first network output for the second network input without providing the first network input as input to the larger neural network.