Entropy-based early stopping for draft tree-based speculative decoding in generative machine learning models

Abstract

Certain aspects of the present disclosure provide techniques and apparatus for machine learning. In an example method, a set of leaf tokens of a draft tree comprising a plurality of tokens is generated using a secondary generative machine learning model associated with a primary generative machine learning model. A set of residual probability distributions for the set of leaf tokens is generated, and a set of entropies is computed based on the set of residual probability distributions. A set of drafting decisions for the set of leaf tokens is generated based on the set of entropies. The secondary generative machine learning model is exited from based on the set of drafting decisions and one or more stopping criteria for generation of a new token using the primary generative machine learning model.

Images