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

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.