Machine learning model processing based on perplexity
A method for operating a machine learning model is presented. The machine learning model includes a plurality of sequential transformer blocks. The method comprises receiving input data at a transformer block and processing the input data via a mixture of experts layer. At an auxiliary classifier, a measure of perplexity of the processed input data is determined. Based on the determined measure of perplexity, one or more experts in a downstream transformer block that will subsequently process the input data are indicated. Weight matrices are then fetched for the indicated one or more experts.
1 . A method for operating a machine learning model including a plurality of sequential transformer blocks, comprising:
receiving input data at a transformer block;
processing at least a portion of the input data via a mixture of experts layer;
at an auxiliary classifier, determining a measure of perplexity of the processed input data;
based on the determined measure of perplexity, indicating one or more experts of a group of experts in a downstream transformer block that will subsequently process at least a portion of the processed input data;
fetching sparsified weight matrices for the indicated one or more experts of the group of experts to generate one or more sparsified experts, wherein a level of sparsity for the sparsified weight matrices is based on the determined measure of perplexity; and
further processing at least a portion of the processed input data via the one or more sparsified experts to generate a prediction.
2 . The method of claim 1 , wherein the downstream transformer block is a next transformer block.
3 . The method of claim 1 , further comprising:
indicating that the processed input data is likely to bypass one or more transformer blocks of the plurality of sequential transformer blocks.
4 . The method of claim 1 , further comprising:
performing a top k selection to select one of the one or more experts of the group of experts that will subsequently process the input data.
5 . The method of claim 1 , wherein the measure of perplexity is determined for a single input data shard.
6 . The method of claim 1 , wherein the measure of perplexity is determined for a group of input data shards.
7 . The method of claim 1 , wherein the measure of perplexity is a loss function.
8 . The method of claim 7 , wherein the loss function is a cross-entropy loss function.
9 . A method for operating a machine learning model, comprising:
at a mixture of experts layer, receiving input data comprising a plurality of input data shards;
sorting the input data shards into batches based on common modalities;
fetching sparsified weight matrices for one or more selected neural network experts of a group of neural network experts, to generate one or more sparsified selected neural network experts, the selected neural network experts trained in modalities represented in the batches;
scheduling each batch for processing by one or more sparsified selected neural network experts trained in a relevant modality; and
processing at least a portion of each batch via respective sparsified selected neural network experts to generate a prediction.
10 . The method of claim 9 , further comprising:
maintaining the fetched sparsified weight matrices for each sparsified selected neural network expert at a node based on a relevant batch processing schedule.
11 . The method of claim 10 , further comprising:
unloading the fetched sparsified weight matrices from the node following processing of a batch; and
fetching sparsified weights for a different selected neural network expert to be loaded onto the node.
12 . The method of claim 9 , wherein scheduling each batch for processing by one or more sparsified selected neural network experts trained in a relevant modality is performed by a reinforcement learning agent.
13 . The method of claim 12 , wherein the reinforcement learning agent is trained in load-balancing.
14 . 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, comprising:
a plurality of sequential transformer blocks configured to receive input data, each transformer block comprising:
a mixture of experts layer configured to process the input data; and
an auxiliary classifier configured to determining a measure of perplexity of the processed input data; and
wherein the one or more processors are configured to:
based on the determined measure of perplexity, indicate one or more experts of a group of experts in a downstream transformer block that will subsequently process at least a portion of the processed input data;
fetch sparsified weight matrices for the indicated one or more experts of the group of experts to generate one or more sparsified experts, wherein a level of sparsity for the sparsified weight matrices is based on the determined measure of perplexity; and
further process at least a portion of the processed input data via the one or more sparsified experts to generate a prediction.
15 . The computing system of claim 14 , wherein the one or more processors are further configured to:
indicate that the processed input data is likely to bypass one or more transformer blocks of the plurality of sequential transformer blocks.
16 . The computing system of claim 14 , wherein the measure of perplexity is determined for a group of input data shards.
17 . The computing system of claim 14 , wherein the measure of perplexity is a loss function.
18 . The method of claim 1 , wherein sparsified weight matrices fetched for experts indicated for processing processed input data with a low level of perplexity have a greater level of sparsity as compared to sparsified weight matrices fetched for experts indicated for processing processed input data with a high level of perplexity.