Decoder architecture for early-stop automated retrieval augmented generation
Aspects of the present disclosure relate to a decoder model architecture that is optimized for retrieval augmented generation (RAG) applications. Embodiments include a generative machine learning model that comprises one or more hidden decoder layers and an output layer. In certain embodiments, the output layer includes a stop head layer configured to generate, based on a hidden state output generated by a given hidden layer of the one or more hidden decoder layers, a stop output indicating a likelihood of performing a subsequent context retrieval. Some embodiments provide that the output layer includes a softmax head layer parallel to the stop head layer. In certain embodiments, each of the one or more hidden decoder layers and the output layer comprise a plurality of neurons organized in an array, wherein each neuron comprises a register, a microprocessor, and at least one input. The neurons may be connected using synaptic circuitry.
1 . A generative machine learning system,
comprising: one or more processors; and
a memory comprising instructions that, when executed by the one or more processors, cause the one or more processors to implement:
a generative machine learning model including: one or more hidden decoder layers; and an output layer comprising:
a stop head layer comprising a plurality of weights organized in a neuron array, wherein the weights are trained so that the stop head layer generates, based on a hidden state output generated by a given hidden layer of the one or more hidden decoder layers, a stop output indicating a likelihood of performing a subsequent context retrieval, and wherein the plurality of weights comprise a smaller number of weights than a number of weights of the one or more hidden decoder layers; and
a softmax head layer parallel to the stop head layer, wherein the softmax head layer generates tokens of a final response based on the hidden state output in response to the stop output not exceeding a threshold; and
a retrieval component configured to perform the subsequent context retrieval based on the stop output exceeding the threshold.
2 . The generative machine learning system of claim 1 , wherein:
each of the one or more hidden decoder layers and the output layer comprise a plurality of neurons organized in an array, wherein each neuron comprises a register, a microprocessor, and at least one input; and
each neuron is connected to at least one other neuron via one of a plurality of synaptic circuits, each synaptic circuit including a memory for storing a synaptic weight.
3 . The generative machine learning system of claim 1 , wherein configuring the stop head layer to generate the stop output comprises training the stop head layer through a supervised learning process involving training data that comprises:
a training stop output generated by the given hidden layer based on an input query and a context; and
a ground truth label that indicates either:
additional context is required to generate a response to the input query; or
no further context is required to generate a response to the input query.
4 . The generative machine learning system of claim 3 , wherein the supervised learning process comprises:
generating a training stop output based on processing the training stop output through the stop head layer; and
adjusting parameters of the stop head layer based on a variance between the training stop output and the ground truth label.
5 . The generative machine learning system of claim 4 , wherein the variance is calculated based on cross-entropy loss.
6 . The generative machine learning system of claim 1 , wherein the generative machine learning model is configured such that the stop head layer is activated at a conclusion of each reasoning step.
7 . The generative machine learning system of claim 1 , further comprising a user interface that is configured to display a final output to a user based on a particular stop output failing to exceed the threshold.
8 . A generative machine learning system,
comprising: one or more processors; and
a memory comprising instructions that, when executed by the one or more processors, cause the one or more processors to:
retrieve a context based on an input query;
generate, via a generative machine learning model, a hidden state output based on the context and the input query;
generate, via a stop head layer of the generative machine learning model based on the hidden state output, a stop output indicating a likelihood of performing a subsequent context retrieval, wherein the stop head layer comprises a plurality of weights organized in a neuron array and wherein the weights are trained to generate the stop output, and wherein the plurality of weights comprise a smaller number of weights than a number of weights of one or more hidden decoder layers of the generative machine learning model; and
either:
perform, based on the stop output exceeding a threshold, the subsequent context retrieval and use the generative machine learning model to generate a subsequent hidden state output based on additional context retrieved in the subsequent context retrieval; or
generate, via an output layer comprising: the stop head layer, and a softmax head layer parallel to the stop head layer, a final output for the input query based on the hidden state output in response to the stop output not exceeding the threshold.
9 . The generative machine learning system of claim 8 , wherein: each of the one or more hidden decoder layers and the output layer comprise a plurality of neurons organized in an array, wherein each neuron comprises a register, a microprocessor, and at least one input; and each neuron is connected to at least one other neuron via one of a plurality of synaptic circuits, each synaptic circuit including a memory for storing a synaptic weight.
10 . The generative machine learning system of claim 8 , wherein the stop head layer to generate the stop output comprises training the stop head layer through a supervised learning process involving training data that comprises:
a training stop output generated based on an input query and a context; and a ground truth label that indicates either:
additional context is required to generate a response to the input query; or no further context is required to generate a response to the input query.
11 . The generative machine learning system of claim 10 , wherein the supervised learning process comprises:
generating a training stop output based on processing the training stop output through the stop head layer; and
adjusting parameters of the stop head layer based on a variance between the training stop output and the ground truth label.
12 . The generative machine learning system of claim 11 , wherein the variance is calculated based on cross-entropy loss.
13 . The generative machine learning system of claim 8 , wherein the generative machine learning model is configured such that the stop head layer is activated at a conclusion of each reasoning step.
14 . The generative machine learning system of claim 8 , further comprising a user interface that is configured to display the final output to a user.
15 . A method,
comprising:
retrieving a context based on an input query;
generating, via a generative machine learning model, a hidden state output based on the context and the input query;
generating, via a stop head layer of the generative machine learning model based on the hidden state output, a stop output indicating a likelihood of performing a subsequent context retrieval, wherein the stop head layer comprises a plurality of weights organized in a neuron array and wherein the weights are trained to generate the stop output, and wherein the plurality of weights comprise a smaller number of weights than a number of weights of one or more hidden decoder layers of the generative machine learning model; and
either:
performing, based on the stop output exceeding a threshold, the subsequent context retrieval and using the generative machine learning model to generate a subsequent hidden state output based on additional context retrieved in the subsequent context retrieval; or
generating, via an output layer comprising: the stop head layer, and a softmax head layer parallel to the stop head layer, final output for the input query based on the hidden state output in response to the stop output not exceeding the threshold.
16 . The method of claim 15 , wherein: each of the one or more hidden decoder layers and the output layer comprise a plurality of neurons organized in an array, wherein each neuron comprises a register, a microprocessor, and at least one input; and each neuron is connected to at least one other neuron via one of a plurality of synaptic circuits, each synaptic circuit including a memory for storing a synaptic weight.
17 . The method of claim 15 , further comprising configuring the stop head layer to generate the stop output by training the stop head layer through a supervised learning process involving training data that comprises:
a training stop output generated based on an input query and a context; and a ground truth label that indicates either:
additional context is required to generate a response to the input query; or no further context is required to generate a response to the input query.
18 . The method of claim 17 , wherein the supervised learning process comprises:
generating a training stop output based on processing the training stop output through the stop head layer; and
adjusting parameters of the stop head layer based on a variance between the training stop output and the ground truth label.