Hybrid speculative decoding in machine learning models

A processing system for machine learning, comprising:

• one or more memories comprising processor-executable instructions;
• and
• one or more processors coupled to the one or more memories and configured to execute the processor-executable instructions and cause the processing system to:
• receive, at a primary generative machine learning model, a first set of tokens generated using a secondary generative machine learning model associated with the primary generative machine learning model;
• identify a first subset of tokens, from the first set of tokens, marked for verification;
• process the first set of tokens using a first subset of layers, of a plurality of layers of the primary generative machine learning model, to generate a first set of hidden tensors;
• select a first subset of the first set of hidden tensors corresponding to the first subset of tokens;
• process the first subset of hidden tensors using a second subset of layers, of the plurality of layers, to generate a set of target logits for the first subset of tokens;
• and
• verify one or more tokens in the first subset of tokens based on the set of target logits.

A processor-implemented method for machine learning, comprising:

• generating, using a secondary generative machine learning model associated with a primary generative machine learning model, a first set of tokens;
• selecting a first subset of tokens, from the first set of tokens, for verification by the primary generative machine learning model;
• sending the first subset of tokens to the primary generative machine learning model for verification;
• and
• generating a machine learning model output comprising one or more verified tokens from the first subset of tokens and one or more remaining tokens from the first set of tokens.