US 12,755,097 B2Granted

Few-shot training of transformer language model from a subset of prompt completion

AssigneeSambaNova Systems, Inc.

InventorsZoltan Csaki, Bo Li, Urmish Ajit Thakker, Venkat Krishna Srinivasan

Life of this patent

Estimating the term…

Loading patent…

Abstract

Embodiments described herein provide systems and techniques for training large language models. In one aspect, a process for performing in-context few-shot training for a transformer-based language model is disclosed. This process may begin by receiving the transformer-based language model having a context window of a predetermined size, as well as a training dataset comprising a set of prompt/completion examples. The process then constructs a training sequence based on the training dataset. Next, the process performs a single forward pass using the training sequence as input. The process subsequently performs a set of backward passes from a subset of examples in the training sequence, wherein each backward pass is conditioned on a selected subset of prompt/completion examples in the training sequence. Performing multiple backward passes conditioned on selected prompt/completion examples improves the language model's ability to generate a higher quality conditional probability distribution for next token without incurring additional training time.

Claims16

Dependent claims

A method of performing in-context few-shot training for a transformer-based language model, comprising:

receiving the transformer-based language model having a context window of a predetermined size;
receiving a training dataset comprising a set of K prompt/completion examples;
constructing a first training sequence based on the training dataset;
performing a single forward pass using the first training sequence as input;
and performing a set of backward passes from a subset of examples in the first training sequence, wherein each backward pass is conditioned on a selected subset of prompt/completion examples in the first training sequence, wherein performing multiple backward passes conditioned on the selected prompt/completion examples in the first training sequence improves the language model's ability to generate a higher quality conditional probability distribution for a next token without incurring additional training time, wherein constructing the first training sequence based on the training dataset includes:
selecting from the set of K prompt/completion examples, a set of N prompt/completion examples, wherein [N £ K] N≤K and the combined size of the selected N prompt/completion examples is equal to or less than the predetermined size of the context window;
and concatenating the selected set of N prompt/completion examples to form the first training sequence, and wherein while performing the single forward pass, the method further comprises:
specifying a set of loss masks within the training sequence for the set of backward passes, wherein each loss mask in the set of loss masks is a unique example in the set of N prompt/completion examples to begin a backward pass;
and specifying a set of attention masks for the set of loss masks, wherein each attention mask in the set of attention masks specifies, during a backward pass from a corresponding loss mask, a subset of prompt/completion examples within the training sequence for the backward pass to be conditioned on.

6 more elements

An apparatus for performing in-context few-shot training for a transformer-based language model, the apparatus comprising:

one or more processors;
a memory coupled to the one or more processors, the memory storing instructions that, when executed by the one or more processors, cause the apparatus to;
receive the transformer-based language model having a context window of a predetermined size;
receive a training dataset comprising a set of K prompt/completion examples;
construct a first training sequence based on the training dataset;
perform a single forward pass using the first training sequence as input;
and perform a set of backward passes from a subset of examples in the first training sequence, wherein each backward pass is conditioned on a selected subset of prompt/completion examples in the first training sequence, wherein performing multiple backward passes conditioned on the selected prompt/completion examples in the first training sequence improves the language model's ability to generate a higher quality conditional probability distribution for a next token without incurring additional training time, wherein constructing the first training sequence based on the training dataset includes:
selecting from the set of K prompt/completion examples, a set of N prompt/completion examples, wherein [N £ K] N≤K and the combined size of the selected N prompt/completion examples is equal to or less than the predetermined size of the context window;
and concatenating the selected set of N prompt/completion examples to form the first training sequence, and wherein while performing the single forward pass, the method further comprises:
specifying a set of loss masks within the training sequence for the set of backward passes, wherein each loss mask in the set of loss masks is a unique example in the set of N prompt/completion examples to begin a backward pass;
and specifying a set of attention masks for the set of loss masks, wherein each attention mask in the set of attention masks specifies, during a backward pass from a corresponding loss mask, a subset of prompt/completion examples within the training sequence for the backward pass to be conditioned on.

8 more elements

Dependencies are inferred from the claim text.