Memory-efficient draft machine learning model

A method of using a machine learning model to generate tokens, the method comprising:

• processing, using a linear layer, an embedding generated from a first output token and input features to generate first features, wherein the first output token is generated by a previous iteration of a token predictor and wherein the input features are generated by a previous iteration of a decoding layer;
• processing, using the decoding layer, the first features to generate second features having first dimensions;
• processing, using a down-projection layer, the second features to generate third features having second dimensions smaller than the first dimensions;
• and
• generating, using the token predictor and the third features, a second output token.

A non-transitory computer-readable medium having stored thereon instructions that, when executed by at least one processor, cause the at least one processor to:

• process, using a linear layer of a machine learning model, an embedding generated from a first output token and input features to generate first features, wherein the first output token is generated by a previous iteration of a token predictor and wherein the input features are generated by a previous iteration of a decoding layer;
• process, using the decoding layer, the first features to generate second features having first dimensions;
• process, using a down-projection layer, the second features to generate third features having second dimensions smaller than the first dimensions;
• and
• generate, using the token predictor and the third features, a second output token.