Augmenting attention-based neural networks to selectively attend to past inputs
Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for performing a machine learning task on a network input that is a sequence to generate a network output. In one aspect, one of the methods includes, for each particular sequence of layer inputs: for each attention layer in the neural network: maintaining episodic memory data; maintaining compressed memory data; receiving a layer input to be processed by the attention layer; and applying an attention mechanism over (i) the compressed representation in the compressed memory data for the layer, (ii) the hidden states in the episodic memory data for the layer, and (iii) the respective hidden state at each of the plurality of input positions in the particular network input to generate a respective activation for each input position in the layer input.
1. A computer-implemented method comprising:
receiving a sequence of network inputs, wherein each network input is an input sequence that comprises a plurality of input positions;
processing the sequence of network inputs using a neural network to generate a network output, wherein the neural network comprises an attention layer associated with a short-term, episodic memory and a long-term, compressed memory, and wherein the processing comprises, for a particular network input in the sequence of network inputs:
updating the short-term, episodic memory to include new data, the new data comprising respective hidden states that were processed by the attention layer for a plurality of input positions in the particular network input;
determining a compression rate for a portion of existing data stored in the short-term, episodic memory, the portion of existing data comprising respective hidden states that were processed by the attention layer for a plurality of input positions in each of one or more previous network inputs that precede the particular network input in the sequence of network inputs;
compressing, in accordance with the compression rate, the portion of existing data to generate a compressed representation for the portion of existing data; and
updating the long-term, compressed memory to include the compressed representation for the portion of existing data.
2. The method of claim 1 , further comprising deleting the portion of existing data from the short-term, episodic memory.
3. The method of claim 1 , wherein the processing comprises, for each input position in the plurality of input positions included in the particular network input:
applying an attention mechanism over (i) data stored in the long-term, compressed memory associated with the attention layer, (ii) data stored in the short-term, episodic memory associated with the attention layer, and (iii) respective hidden states that were processed by the attention layer for one or more input positions that precede the input position in the plurality of input positions, to generate an activation for the input position.
4. The method of claim 1 , wherein compressing, in accordance with the compression rate, the portion of existing data comprises:
applying a max pooling function to the portion of existing data with a stride equal to the compression rate.
5. The method of claim 1 , wherein compressing, in accordance with the compression rate, the portion of existing data comprises:
applying a mean pooling function to the portion of existing data with a stride equal to the compression rate.
6. The method of claim 1 , wherein compressing, in accordance with the compression rate, the portion of existing data comprises:
determining a kernel size for a 1D convolution function; and
applying the 1D convolution function to the portion of existing data with a stride equal to the compression rate and a kernel size equal to the determined kernel size.
7. The method of claim 1 , wherein compressing, in accordance with the compression rate, the portion of existing data comprises:
applying a multi-layer dilated 1D convolution function to the portion of existing data.
8. The method of claim 1 , wherein the one or more previous network inputs further precede one or more other network inputs that precede the particular network input in the sequence of network inputs.
9. The method of claim 8 , wherein a number of the one or more other network inputs is dependent on a predetermined size of the short-term, episodic memory.
10. The method of claim 1 , wherein the processing comprises:
processing activations generated by the attention layer for the plurality of input positions using one or more fully connected layers and one or more layer normalization layers included in the neural network.
11. The method of claim 1 , wherein the sequence of network inputs comprises data derived from a sequence of multiple articles from a text document.
12. The method of claim 1 , wherein the sequence of network inputs comprises data derived from a sequence of multiple observations generated while an agent interacts with an environment.
13. A system comprising:
one or more computers; and
one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations comprising:
receiving a sequence of network inputs, wherein each network input is an input sequence that comprises a plurality of input positions;
processing the sequence of network inputs using a neural network to generate a network output, wherein the neural network comprises an attention layer associated with a short-term, episodic memory and a long-term, compressed memory, and wherein the processing comprises, for a particular network input in the sequence of network inputs:
updating the short-term, episodic memory to include new data, the new data comprising respective hidden states that were processed by the attention layer for a plurality of input positions in the particular network input;
determining a compression rate for a portion of existing data stored in the short-term, episodic memory, the portion of existing data comprising respective hidden states that were processed by the attention layer for a plurality of input positions in each of one or more previous network inputs that precede the particular network input in the sequence of network inputs;
compressing, in accordance with the compression rate, the portion of existing data to generate a compressed representation for the portion of existing data; and
updating the long-term, compressed memory to include the compressed representation for the portion of existing data.
14. The system of claim 13 , wherein the processing comprises:
processing activations generated by the attention layer for the plurality of input positions using one or more fully connected layers and one or more layer normalization layers included in the neural network.
15. The system of claim 13 , wherein the sequence of network inputs comprises data derived from a sequence of multiple articles from a text document.
16. The system of claim 13 , wherein the sequence of network inputs comprises data derived from a sequence of multiple observations generated while an agent interacts with an environment.
17. One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:
receiving a sequence of network inputs, wherein each network input is an input sequence that comprises a plurality of input positions;
processing the sequence of network inputs using a neural network to generate a network output, wherein the neural network comprises an attention layer associated with a short-term, episodic memory and a long-term, compressed memory, and wherein the processing comprises, for a particular network input in the sequence of network inputs:
updating the short-term, episodic memory to include new data, the new data comprising respective hidden states that were processed by the attention layer for a plurality of input positions in the particular network input;
determining a compression rate for a portion of existing data stored in the short-term, episodic memory, the portion of existing data comprising respective hidden states that were processed by the attention layer for a plurality of input positions in each of one or more previous network inputs that precede the particular network input in the sequence of network inputs;
compressing, in accordance with the compression rate, the portion of existing data to generate a compressed representation for the portion of existing data; and
updating the long-term, compressed memory to include the compressed representation for the portion of existing data.
18. The non-transitory computer-readable storage media of claim 17 , wherein the processing comprises:
processing activations generated by the attention layer for the plurality of input positions using one or more fully connected layers and one or more layer normalization layers included in the neural network.
19. The non-transitory computer-readable storage media of claim 17 , wherein the sequence of network inputs comprises data derived from a sequence of multiple articles from a text document.
20. The non-transitory computer-readable storage media of claim 17 , wherein the sequence of network inputs comprises data derived from a sequence of multiple observations generated while an agent interacts with an environment.