Mixture of experts models with sparsified weights
A method is presented for operating a machine learning model including one or more mixture of experts layers. The method comprises receiving one or more input data shards at a routing gate network for a mixture of experts layer comprising a plurality of neural network experts. One or more neural network experts in the mixture of experts layer is designated layer to evaluate each input data shard. For each designated neural network expert, a weight matrix is retrieved having a predetermined sparsity to generate a sparsified designated neural network expert. Each input data shard is evaluated with a respective sparsified designated neural network expert.
1 . A method for operating a machine learning model including one or more mixture of experts layers, the method comprising:
at a routing gate network for a mixture of experts layer that includes a plurality of neural network experts, receiving one or more input data shards related to a predictive modeling problem;
applying coarse-grained sparsity to the mixture of experts layer to reduce a global sparsity level of the mixture of experts layer by designating one or more neural network experts as a subset of the plurality of neural network experts in the mixture of experts layer to evaluate each input data shard;
applying fine-grained sparsity to the designated neural network experts to further reduce the global sparsity level of the mixture of experts layer by, for each designated neural network expert, retrieving a weight matrix having a predetermined sparsity to generate a sparsified designated neural network expert that has a set of one or more parameters including fewer parameters than the neural network expert from which the sparsified designated neural network expert was generated;
evaluating each input data shard with a respective sparsified designated neural network expert to generate a prediction;
outputting a unified prediction at an output gate of the mixture of experts layer that is based, at least in part, on the prediction generated by the sparsified designated neural network expert for each input data shard; and
outputting a solution to the predictive modeling problem from the machine learning model based, at least in part, on the unified prediction.
2 . The method of claim 1 , wherein the predetermined sparsity is configured as unstructured sparsity.
3 . The method of claim 1 , wherein the predetermined sparsity is configured as balanced N of M sparsity, and wherein each neural network expert is implemented using hardware configured to implement balanced N of M sparsity.
4 . The method of claim 3 , wherein the balanced N of M sparsity has a block size of M≥8.
5 . The method of claim 1 , wherein the input data shards include image patches.
6 . The method of claim 1 , wherein the input data shards include sentence tokens.
7 . The method of claim 1 , wherein evaluating the input data shards by neural network experts includes performing inference on the input data shards.
8 . The method of claim 1 , wherein the designation of the one or more neural network experts as the subset of the plurality of neural network experts is based on a top-K function, and wherein the predetermined sparsity is based at least on a ratio of the top-K function to the total number of neural network experts.
9 . The method of claim 1 , wherein the predetermined sparsity is based at least on a level of expertise of the designated neural network expert.
10 . The method of claim 1 , wherein a level of sparsity of the predetermined sparsity is based on a quantity of the designated neural network experts.
11 . The method of claim 1 , wherein the received input data shards are evaluated by two or more sparsified designated neural network experts, and wherein the unified prediction is output based on the predictions generated by each of the sparsified designated neural network experts.
12 . A computing system, comprising:
one or more processors; and
a storage machine having instructions stored thereon executable by the one or more processors to instantiate a machine learning model including one or more mixture of experts layers, each mixture of experts layer comprising:
a plurality of neural network experts;
an output gate; and
a routing gate network configured to:
receive one or more input data shards related to a predictive modeling problem; and
apply coarse-grained sparsity to the mixture of experts layer to reduce a global sparsity level of the mixture of experts layer by designating one or more neural network experts as a subset of the plurality of neural network experts in the mixture of experts layer to evaluate each input data shard; wherein,
for each mixture of experts layer, the one or more processors through execution of the instructions are configured to:
apply fine-grained sparsity to the designated neural network experts to further reduce the global sparsity level of the mixture of experts layer by, for each neural network expert designated to evaluate an input data shard received by the mixture of experts layer, retrieving a weight matrix having a predetermined sparsity to generate a sparsified designated neural network expert that has a set of one or more parameters including fewer parameters than the neural network expert from which the sparsified designated neural network expert was generated;
supervise evaluation of each input data shard received by the mixture of experts layer to generate a prediction with the sparsified designated neural network expert for the input data shard;
output a unified prediction at the output gate of the mixture of experts layer that is based, at least in part, on the prediction generated by the sparsified designated neural network expert for the input data shard; and wherein
the one or more processors through execution of the instructions are further configured to output a solution to the predictive modeling problem from the machine learning model based, at least in part, on the unified prediction.
13 . The computing system of claim 12 , wherein the predetermined sparsity is configured as balanced N of M sparsity, and wherein each neural network expert is implemented using hardware configured to implement balanced N of M sparsity.
14 . The computing system of claim 13 , wherein the balanced N of M sparsity has a block size of M≥8.
15 . The computing system of claim 12 , wherein the input data shards include image patches.
16 . The computing system of claim 12 , wherein the input data shards include sentence tokens.
17 . The computing system of claim 12 , wherein evaluating the input data shards by neural network experts includes performing inference on the input data shards.
18 . The computing system of claim 12 , wherein the designation of the one or more neural network experts as the subset of the plurality of neural network experts is based on a top-K function, and wherein the predetermined sparsity is based at least on a ratio of the top-K function to the total number of neural network experts.
19 . The computing system of claim 12 , wherein the predetermined sparsity is based at least on a level of expertise of the designated neural network expert.
20 . The computing system of claim 12 , wherein a level of sparsity of the predetermined sparsity is based on a depth of the mixture of experts layer among the one or more mixture of experts layers.