Dynamic cache management in beam search
Systems and methods for dynamically modifying a cache associated with a neural network model of a natural language generator are described. In examples, a neural network model employs a beam search algorithm at a decoder when decoding output and generating predicted output candidates. The decoder utilizes caching techniques to improve a speed at which the neural network operations. When an amount of memory utilized by one or more caches of the neural network model is determined to exceed a threshold memory size, a layer-specific portion of a cache associated with a layer of the neural network model is identified. The identified layer-specific portion of the cache can be deleted when the amount of memory utilized by the cache of the neural network model exceeds the threshold memory size. In examples, data in the cache is deduplicated and/or deleted.
1 . A system comprising:
a cache including a plurality of layer-specific caches storing encoded inputs generated by an encoder of a neural network model, the encoded inputs used by a decoder of the neural network model to generate output candidates, wherein the decoder includes a plurality of layers, where a first layer-specific cache of the plurality of layer-specific caches corresponds to a first layer of the plurality of layers of the decoder;
a processor; and
memory, including instructions that, as a result of being executed by the processor, cause the processor to:
determine that an amount of memory utilized by the cache exceeds a threshold memory size, wherein the amount of memory allocated to the cache is determined based on batch input processing and the plurality of layer-specific caches;
determine the first layer-specific cache includes a first portion of the encoded inputs that are not associated with another operation of the decoder;
modify the cache by discarding the first portion of the encoded inputs stored in the first layer-specific cache such that the cache is resized based on discarding the first portion of the encoded inputs; and
decode a second portion of the encoded inputs utilizing a beam search algorithm using the cache.
2 . The system of claim 1 , further comprising instructions, that, as a result of being executed by the processor, cause the processor to:
determine that the amount of memory utilized by the cache is less than the threshold memory size;
determine a second layer of the plurality of layers of the decoder that is not associated with data in the cache; and
generate additional data for the second layer, wherein the additional data is stored in a second layer-specific cache of the plurality of layer-specific caches corresponding to the second layer.
3 . The system of claim 2 , wherein the additional data is generated by deduplicating data stored in the second layer-specific cache.
4 . The system of claim 1 , wherein discarding the first portion of the encoded inputs further comprises deduplicating data stored in the first layer-specific cache.
5 . The system of claim 4 , further comprising instructions, that, as a result of being executed by the processor, cause the processor to:
receive data associated with an entry in the first layer-specific cache for a first beam in the beam search algorithm; and
cause the first layer to generate a first output candidate utilizing the entry.
6 . The system of claim 4 , wherein a number of beam-specific portions of the cache that are associated with layers of the plurality of layers is less than a beamwidth employed by the beam search algorithm.
7 . The system of claim 1 , wherein the first layer-specific cache is associated with a lowest layer of the decoder.
8 . A method comprising:
receiving an input at an encoder of a neural network model;
encoding the input at the encoder of the neural network model to generate an encoded input;
causing a decoder of the neural network model to generate output candidates based on the encoded input, wherein the decoder includes a plurality of layers utilizing a plurality of layer-specific caches of a cache that stores the encoded input;
obtaining a first data associated with an entry in a first layer-specific cache for a first beam in a beam search algorithm;
modifying the first layer-specific cache by discarding a second data stored in the first layer-specific cache as a result of an amount of memory utilized by the cache exceeding a threshold memory size, wherein the cache is dynamically resized as a result of discarding a second data stored in the first layer-specific cache;
causing a first layer of the decoder to generate a first output candidate based on the entry in the first layer-specific cache;
obtaining a second data associated with the entry in the first layer-specific cache for a second beam in the beam search algorithm; and
causing a first layer of the decoder to generate a second output candidate utilizing the entry in the first layer-specific cache.
9 . The method of claim 8 , wherein a number of beam-specific portions of the cache is less than a beamwidth parameter used by the beam search algorithm.
10 . The method of claim 8 , further comprising:
determining that the amount of memory utilized by the cache is less than a threshold memory size;
identifying a layer of the decoder that is not associated with data stored in the cache; and
generating a third data for the layer, wherein the third data is stored in a third layer-specific cache that is associated with the layer of the decoder.
11 . The method of claim 8 , wherein discarding the second data stored further comprises deduplicating data stored in the first layer-specific cache.
12 . The method of claim 8 , wherein the first layer-specific cache is associated with a lowest layer of the neural network model.
13 . A system comprising:
a cache including a first layer-specific cache storing encoded inputs generated by an encoder of a neural network model, the encoded inputs used by a decoder of the neural network model to generate output candidates, wherein the decoder includes a plurality of layers, where the first layer-specific cache corresponds to a first layer of the plurality of layers of the decoder and an amount of memory allocated to the cache is determined based on a tradeoff between memory allocated for batch input processing and memory allocated for caching layer-specific data;
a processor; and
memory including instructions that, as a result of being executed by the processor, cause the processor to:
receive an input at the encoder of the neural network model;
cause the encoder to encode the input to generate an encoded input; and
cause the decoder to generate output candidates based on the encoded input stored in the first layer-specific cache, where the first layer-specific cache includes deduplicated data, wherein the decoder generates the output candidates by:
obtaining data associated with an entry in the first layer-specific cache for a first beam in a beam search algorithm;
generating a first output candidate utilizing the entry in the first layer-specific cache;
modifying the cache by discarding a first portion of data stored in the first layer-specific cache in response to an amount of memory utilized by the cache exceeding a threshold memory size;
obtaining additional data associated with the entry for a second beam in the beam search algorithm; and
generating a second output candidate utilizing the entry.
14 . The system of claim 13 , wherein a number of beams of the beam search algorithm is less than a beamwidth parameter used by the beam search algorithm.
15 . The system of claim 13 , further comprising instructions, that, as a result of being executed by the processor, cause the processor to:
determine that the amount of memory utilized by the cache is less than a threshold memory size;
identify a second layer of the decoder that is not associated with data in the cache; and
generates second data for the second layer, wherein the second data is stored in a second layer-specific cache that is associated with the second layer.
16 . The system of claim 13 , wherein discarding the first portion of data includes deduplicating data stored in the first layer-specific cache.
17 . The system of claim 13 , wherein the first layer-specific cache is associated with a lowest layer of the decoder.
18 . The system of claim 1 , wherein the first portion of the encoded inputs that are not associated with another operation of the decoder are determined based at least in part on a flag indicating that data in the first layer-specific cache is not needed.
19 . The method of claim 8 , wherein discarding the second data is based on a flag indicating that the second data in the first layer-specific cache is not needed for an operation of the decoder.
20 . The system of claim 13 , wherein discarding the first portion of data is based on a flag indicating that data in the first layer-specific cache is not needed.