Domain-aware large language model governance
Techniques are described herein for a method of decreasing the likelihood of out-of-domain LLM responses. The method includes determining, by a block of a LLM, a representation of the text input. The method further includes determining a set of coefficients based at least on a reconstruction of the text input using a dictionary and the representation of the text input. The method further includes performing a sparsity check using the set of coefficients. The method further includes generating a response to the text input based at least on the sparsity check.
1 . A method to constrain responses being generated by a large language model (LLM), the method comprising:
generating, by a prompt manager, an input for the LLM, wherein the input comprises a plurality of tokens;
generating, by the LLM being executed on a processing device, a response to the input, wherein the generating includes:
determining, by a block of the LLM, an output that includes a representation of at least one of the plurality of tokens;
determining a set of coefficients based at least on the output and a template predictor, wherein the template predictor was previously determined to predict templates in a dictionary responsive to outputs from the block, wherein the dictionary defines a previously detected sub-space representing outputs of the block that have a characteristic, wherein the characteristic represents one or more portions of a specific domain or alignment with a policy;
performing a sparsity check using the set of coefficients; and
constraining the response being generated by the LLM when it is determined the response being generated is not within the previously detected subspace based at least in part on the sparsity check; and
causing the response to be presented to a user as at least part of a subsequent turn in a conversation in which the user is actively engaged.
2 . The method of claim 1 , wherein the performing comprises:
comparing a number of non-zero elements in the set of coefficients to a sparsity pattern.
3 . The method of claim 2 , further comprising:
determining that the response being generated is within the previously detected subspace based on the number of non-zero elements in the set of coefficients satisfying the sparsity pattern.
4 . The method of claim 2 , further comprising:
determining that the response being generated is not within the previously detected subspace based on the number of non-zero elements in the set of coefficients not satisfying the sparsity pattern.
5 . The method of claim 4 , further comprising:
tracking a number of failed sparsity checks.
6 . The method of claim 1 , wherein the dictionary is a first dictionary, and the set of coefficients is a first set of coefficients, further comprising:
determining, by a second block of the LLM, a second output that includes a representation of at least one of the plurality of tokens;
determining a second set of coefficients based at least on the second output and a second dictionary; and
performing a second sparsity check using the second set of coefficients, wherein the constraining is also based at least in part on the second sparsity check.
7 . The method of claim 1 , wherein the block of the LLM is at least one of an encoder, a decoder, a transformer, a feed-forward layer, a self-attention layer, or a cross-attention layer.
8 . A non-transitory computer-readable storage medium storing executable instructions, which when executed by a processing device, cause the processing device to perform operations comprising:
generating, by a prompt manager, an input for a large language model (LLM), wherein the input comprises a plurality of tokens;
generating, by the LLM, a response to the input, wherein the generating includes:
determining, by a block of the LLM, an output that includes a representation of at least one of the plurality of tokens;
determining a set of coefficients based at least on the output and a template predictor, wherein the template predictor was previously determined to predict templates in a dictionary responsive to outputs from the block, wherein the dictionary defines a previously detected sub-space representing outputs of the block that have a characteristic, wherein the characteristic represents one or more portions of a specific domain or alignment with a policy;
performing a sparsity check using the set of coefficients; and
constraining the response being generated by the LLM when it is determined the response being generated is not within the previously detected subspace based at least in part on the sparsity check; and
causing the response to be presented to a user as at least part of a subsequent turn in a conversation in which the user is actively engaged.
9 . The non-transitory computer-readable storage medium of claim 8 , wherein the performing comprises:
comparing a number of non-zero elements in the set of coefficients to a sparsity pattern.
10 . The non-transitory computer-readable storage medium of claim 9 , wherein the operations further comprise:
determining that the response being generated is within the previously detected subspace based on the number of non-zero elements in the set of coefficients satisfying the sparsity pattern.
11 . The non-transitory computer-readable storage medium of claim 9 , wherein the operations further comprise:
determining that the response being generated is not within the previously detected subspace based on the number of non-zero elements in the set of coefficients not satisfying the sparsity pattern.
12 . The non-transitory computer-readable storage medium of claim 11 , wherein the operations further comprise:
tracking a number of failed sparsity checks.
13 . The non-transitory computer-readable storage medium of claim 8 , wherein the operations further comprise:
determining, by a second block of the LLM, a second output that includes a second representation of at least one of the plurality of tokens;
determining a second set of coefficients based at least on the second output and a second dictionary, wherein the second dictionary is different from the dictionary; and
performing a second sparsity check using the second set of coefficients, wherein the constraining is also based at least in part on the second sparsity check.
14 . The non-transitory computer-readable storage medium of claim 8 , wherein the block of the LLM is at least one of an encoder, a decoder, a transformer, a feed-forward layer, a self-attention layer, or a cross-attention layer.
15 . A system comprising:
a memory component; and
a processing device coupled to the memory component, the processing device to perform operations comprising:
generating, by a prompt manager, and input for a large language model (LLM), wherein the input comprises a plurality of tokens;
generating, by the LLM, a response to the input, wherein the generating includes:
determining, by a block of the LLM, an output that includes a representation of at least one of the plurality of tokens;
determining a set of coefficients based at least on the output and a template predictor, wherein the template predictor was previously determined to predict templates in a dictionary responsive to outputs from the block, wherein the dictionary defines a previously detected sub-space representing outputs of the block that have a characteristic, wherein the characteristic represents one or more portions of a specific domain or alignment with a policy;
performing a sparsity check using the set of coefficients; and
constraining the response being generated by the LLM when it is determined the response being generated is not within the previously detected subspace based at least in part on the sparsity check; and
causing the response to be presented to a user as at least part of a subsequent turn in a conversation in which the user is actively engaged.
16 . The system of claim 15 , wherein the operations further comprise:
determining, by a second block of the LLM, a second output that includes a second representation of at least one of the plurality of tokens;
determining a second set of coefficients based at least on the second output and a second dictionary, wherein the second dictionary is different from the dictionary; and
performing a second sparsity check using the second set of coefficients, wherein the constraining is also based at least in part on the second sparsity check.
17 . The system of claim 15 , wherein the performing comprises:
comparing a number of non-zero elements in the set of coefficients to a sparsity pattern.
18 . The system of claim 17 , wherein the operations further comprise:
determining that the response being generated is within the previously detected subspace based on the number of non-zero elements in the set of coefficients satisfying the sparsity pattern.
19 . The system of claim 17 , wherein the operations further comprise:
determining that the response being generated is not within the previously detected subspace based on the number of non-zero elements in the set of coefficients not satisfying the sparsity pattern.
20 . The system of claim 15 , wherein the operations further comprise:
tracking a number of failed sparsity checks.
21 . The method of claim 1 , wherein the constraining includes:
generating, based on the coefficients and the dictionary, a reconstruction representing an in-domain replacement for the output from the block; and
replacing the output from the block with the reconstruction as the input to a device next block of the LLM.
22 . The method of claim 1 , wherein the constraining includes causing the response generated by the LLM to be a predetermined response rather than one generated using the blocks of the LLM.
23 . The non-transitory computer-readable storage medium of claim 8 , wherein the constraining includes:
generating, based on the coefficients and the dictionary, a reconstruction representing an in-domain replacement for the output from the block; and
replacing the output from the block with the reconstruction as the input to a device next block of the LLM.
24 . The non-transitory computer-readable storage medium of claim 8 , wherein the constraining includes causing the response generated by the LLM to be a predetermined response rather than one generated using the blocks of the LLM.
25 . The system of claim 15 , wherein the constraining includes:
generating, based on the coefficients and the dictionary, a reconstruction representing an in-domain replacement for the output from the block; and
replacing the output from the block with the reconstruction as the input to a device next block of the LLM.
26 . The system of claim 15 , wherein the constraining includes causing the response generated by the LLM to be a predetermined response rather than one generated using the blocks of the LLM.