Systems and methods for training a neural network model using knowledge from pre-trained large language models
Embodiments described herein provide a training framework for generative NLP models that operate on previously learnt knowledge from pretrained large language models. Specifically, to train an NLP model to generate a response to a user utterance (e.g., “resolve login issue”), document embeddings of support IT documents encoded by a pretrained LLM are fed to an NLP decoder together with a training dialogue (e.g., a dialogue between the chat agent on how to “resolve login issue”). The NLP decoder can thus be trained by a causal language modeling loss computed based on the predicted next token and the ground-truth token from the training dialogue.
1 . A method for training a neural network based natural language processing (NLP) model, comprising:
receiving, at a server and via a communication interface over a network, a training input including tokens of a prior user-agent dialogue;
retrieving, at the server, one or more precomputed document embeddings of one or more knowledge documents from a pretrained large language model (LLM);
generating an augmented training input by combining the one or more precomputed document embeddings and token embeddings corresponding to at least a subset of tokens from the training input in an embedding space, wherein the generating comprises inserting a first token indicating a first knowledge document from the one or more knowledge documents at a start of a first prior system response in the prior user-agent dialogue;
training the neural network based NLP model implemented on one or more hardware processors at the server using the augmented training input, wherein the training comprises:
predicting, by the neural network based NLP model, a document identification in response to the augmented training input,
computing a factual alignment loss by comparing the predicted document identification and the first token, and
updating parameters of the neural network based NLP model based on the factual alignment loss; and
generating, by the trained neural network based NLP model, an agent response that is based on the one or more knowledge documents in response to a user utterance.
2 . The method of claim 1 , wherein the prior user-agent dialogue comprises a prior user utterance and at least one prior system response generated in response to the prior user utterance based on at least one of the one or more knowledge documents.
3 . The method of claim 1 , wherein the one or more precomputed document embeddings are retrieved via an application programming interface (API) at the server, which connects the server to an external server that hosts the pretrained LLM, and
wherein the one or more precomputed document embeddings are generated by the pretrained LLM and stored at the external server prior to training the neural network based natural language processing (NLP) model.
4 . The method of claim 1 , wherein the training the neural network based NLP model using the augmented training input further comprises:
predicting, by the neural network based NLP model, a next token subsequent to the subset of tokens in response to the augmented training input;
computing a causal language modeling loss based on the predicted next token and a ground-truth token from the tokens of a prior user-agent dialogue; and
updating parameters of the neural network based NLP model based on the causal language modeling loss.
5 . The method of claim 1 , further comprising:
generating, by the trained neural network based NLP model, a document identification token at a start of the agent response in response to an user utterance, wherein the document identification token indicates which knowledge document the agent response is based on.
6 . The method of claim 1 , further comprising:
pretraining the neural network based NLP model by:
randomly replacing one or more spans in the training input with one or more corresponding embeddings, thereby resulting in a masked training input,
wherein the one or more corresponding embeddings have a same dimension of the neural network based NLP model; and
training the neural network NLP model using a pre-processed dataset including the masked training input.
7 . The method of claim 6 , wherein at least one span includes a sentence or a paragraph, and wherein the one or more corresponding embeddings are generated by the pretrained LLM.
8 . The method of claim 1 , further comprising:
transforming multiple randomly selected training samples into corresponding embeddings generated by the pretrained LLM; and
training the neural network based NLP model using the corresponding embeddings as in-context examples.
9 . A system for training a neural network based natural language processing (NLP) model, comprising:
a memory that stores the neural network based NLP model and a plurality of processor executable instructions;
a communication interface that receives a training input including tokens of a prior user-agent dialogue; and
one or more hardware processors that read and execute the plurality of processor-executable instructions from the memory to perform operations comprising:
retrieving, at the server, one or more precomputed document embeddings of one or more knowledge documents from a pretrained large language model (LLM);
generating an augmented training input by combining the one or more precomputed document embeddings and token embeddings corresponding to at least a subset of tokens from the training input in an embedding space, wherein the generating comprises inserting a first token indicating a first knowledge document from the one or more knowledge documents at a start of a first prior system response in the prior user-agent dialogue;
training the neural network based NLP model implemented on one or more hardware processors at the server using the augmented training input, wherein the training comprises:
predicting, by the neural network based NLP model, a document identification in response to the augmented training input,
computing a factual alignment loss by comparing the predicted document identification and the first token, and
updating parameters of the neural network based NLP model based on the factual alignment loss; and
generating, by the trained neural network based NLP model, an agent response that is based on the one or more knowledge documents in response to a user utterance.
10 . The system of claim 9 , wherein the prior user-agent dialogue comprises a prior user utterance and at least one prior system response generated in response to the prior user utterance based on at least one of the one or more knowledge documents.
11 . The system of claim 9 , wherein the one or more precomputed document embeddings are retrieved via an application programming interface (API) at the server, which connects the server to an external server that hosts the pretrained LLM, and
wherein the one or more precomputed document embeddings are generated by the pretrained LLM and stored at the external server prior to training the neural network based natural language processing (NLP) model.
12 . The system of claim 9 , wherein the operation of training the neural network based NLP model using the augmented training input further comprises:
predicting, by the neural network based NLP model, a next token subsequent to the subset of tokens in response to the augmented training input;
computing a causal language modeling loss based on the predicted next token and a ground-truth token from the tokens of a prior user-agent dialogue; and
updating parameters of the neural network based NLP model based on the causal language modeling loss.
13 . The system of claim 9 , wherein the operations further comprise:
generating, by the trained neural network based NLP model, a document identification token at a start of the agent response in response to an user utterance, wherein the document identification token indicates which knowledge document the agent response is based on.
14 . The system of claim 9 , wherein the operations further comprise:
pretraining the neural network based NLP model by:
randomly replacing one or more spans in the training input with one or more corresponding embeddings, thereby resulting in a masked training input,
wherein the one or more corresponding embeddings have a same dimension of the neural network based NLP model; and
training the neural network NLP model using a pre-processed dataset including the masked training input.
15 . The system of claim 14 , wherein at least one span includes a sentence or a paragraph, and wherein the one or more corresponding embeddings are generated by the pretrained LLM.
16 . A non-transitory machine-readable medium comprising a plurality of machine-executable instructions which, when executed by one or more processors, are adapted to cause the one or more processors to perform operations comprising:
receiving, at a server and via a communication interface over a network, a training input including tokens of a prior user-agent dialogue;
retrieving, at the server, one or more precomputed document embeddings of one or more knowledge documents from a pretrained large language model (LLM);
generating an augmented training input by combining the one or more precomputed document embeddings and token embeddings corresponding to at least a subset of tokens from the training input in an embedding space, wherein the generating comprises inserting a first token indicating a first knowledge document from the one or more knowledge documents at a start of a first prior system response in the prior user-agent dialogue;
training the neural network based NLP model implemented on one or more hardware processors at the server using the augmented training input, wherein the training comprises:
predicting, by the neural network based NLP model, a document identification in response to the augmented training input,
computing a factual alignment loss by comparing the predicted document identification and the first token, and
updating parameters of the neural network based NLP model based on the factual alignment loss; and
generating, by the trained neural network based NLP model, an agent response that is based on the one or more knowledge documents in response to a user utterance.