Neural networks with relational memory
A system including 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 implement a memory and memory-based neural network is described. The memory is configured to store a respective memory vector at each of a plurality of memory locations in the memory. The memory-based neural network is configured to: at each of a plurality of time steps: receive an input; determine an update to the memory, wherein determining the update comprising applying an attention mechanism over the memory vectors in the memory and the received input; update the memory using the determined update to the memory; and generate an output for the current time step using the updated memory.
1. 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 implement:
a memory configured to store a respective memory vector at each of a plurality of memory locations in the memory; and
a memory-based neural network, wherein the memory-based neural network is configured to:
at each of a plurality of time steps:
receive an input;
determine an update to the memory, wherein determining the update comprises applying an attention mechanism over (i) the memory vectors in the memory and (ii) memory vectors in an appended memory that is generated by appending the received input to the memory, wherein applying the attention mechanism over the memory vectors in the memory and memory vectors in the appended memory comprises:
generating, from a subset of the memory vectors in the memory, a query matrix,
generating, from a subset of the memory vectors in the appended memory, a key matrix,
generating, from the subset of the memory vectors in the appended memory, a value matrix,
generating, from (i) the query matrix generated from the subset of the memory vectors in the memory and (ii) the key matrix generated from the subset of the memory vectors in the appended memory, a weight matrix, and
determining a first update to the memory based on the weight matrix and the value matrix generated from the subset of the memory vectors in the appended memory;
update the memory using the determined update to the memory; and
generate an output for the current time step using the updated memory.
2. The system of claim 1 , wherein the input received at each of the plurality of time step is one of (i) a network input of the memory-based neural network for the current time step or (ii) an intermediate input derived from a network input of the memory-based neural network for the current time step.
3. The system of claim 1 , wherein the memory-based neural network
comprises a plurality of self-attention layers, each self-attention layer corresponding to a respective subset of the memory vectors of the memory, and a respective subset of the memory vectors of the appended memory;
wherein at each time step, applying the attention mechanism over the memory vectors in the memory and memory vectors in the appended memory comprises:
determining the first update to the memory by using one of the plurality of self-attention layers and its respective parameter set, and
determining proposed updates to the memory by using other self-attention layers of the plurality of self-attention layers, wherein each of the other self-attention layers is configured to determine a respective proposed update to the memory by applying an attention mechanism over the respective subset of memory vectors in the memory and the respective subset of memory vectors in the appended memory; and
wherein the memory-based neural network is further configured to, at each time step, combine the first update with the proposed updates to generate the update for the memory.
4. The system of claim 3 , wherein the self-attention layers operate in parallel.
5. The system of claim 3 , wherein combining the first update with the proposed updates comprises column-wise concatenating the first update and the proposed updates to generate the update for the memory.
6. The system of claim 3 ,
wherein at each time step, applying the attention mechanism over the memory vectors in the memory and memory vectors in the appended memory comprises:
generating, by the one of the plurality of self-attention layers, from a subset of the memory vectors in the memory, the query matrix, by:
applying a query linear projection of the subset of the memory vectors in the memory onto a query weight matrix to generate the query matrix,
generating, by the one of the plurality of self-attention layers, from a subset of the memory vectors in the appended memory, the key matrix by:
applying a key linear projection of the subset of memory vectors in the appended memory onto a key weight matrix to generate the key matrix,
generating, by the one of the plurality of self-attention layers, from the subset of the memory vectors in the appended memory, the value matrix by:
applying a value linear projection of the subset of memory vectors in the appended memory onto a value weight matrix to generate the value matrix,
generating, from (i) the query matrix generated from the subset of the memory vectors in the memory and (ii) the key matrix generated from the subset of the memory vectors in the appended memory, a weight matrix by:
multiplying the query matrix and a transpose of the key matrix to determine a first temporary matrix,
dividing each element of the first temporary matrix by a scaling factor to determine a second temporary matrix,
applying a softmax operator on each element of the second temporary matrix to generate the weight matrix, and
determining the first update to the memory based on the weight matrix and the value matrix generated from the subset of the memory vectors in the appended memory by:
multiplying the weight matrix and the respective value matrix to determine the first update to the memory.
7. The system of claim 6 , wherein the scaling factor is a square root of the dimension of key vectors in the key matrix.
8. The system of claim 3 , wherein the memory-based neural network further comprises a feedforward neural network and one or more memory gating layers; and
wherein at each time step, generating the output for the current time step using the updated memory comprises:
processing the updated memory using the feedforward neural network to generate an intermediate output; and
gating the intermediate output using the one or more memory gating layers to generate the output for the current time step.
9. The system of claim 8 , wherein the feedforward neural network is a multilayer perceptron (MLP).
10. The system of claim 3 , wherein the input is an input vector, and wherein appending the input to the memory comprises performing a row-wise concatenation of the memory and the input vector.
11. The system of claim 3 , wherein the input is an input matrix.
12. One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to implement:
a memory configured to store a respective memory vector at each of a plurality of memory locations in the memory; and
a memory-based neural network, wherein the memory-based neural network is configured to:
at each of a plurality of time steps:
receive an input;
determine an update to the memory, wherein determining the update comprises applying an attention mechanism over (i) the memory vectors in the memory and (ii) memory vectors in an appended memory that is generated by appending the received input to the memory, wherein applying the attention mechanism over the memory vectors in the memory and memory vectors in the appended memory comprises:
generating, from a subset of the memory vectors in the memory, a query matrix,
generating, from a subset of the memory vectors in the appended memory, a key matrix,
generating, from the subset of the memory vectors in the appended memory, a value matrix,
generating, from (i) the query matrix generated from the subset of the memory vectors in the memory and (ii) the key matrix generated from the subset of the memory vectors in the appended memory, a weight matrix, and
determining a first update to the memory based on the weight matrix and the value matrix generated from the subset of the memory vectors in the appended memory;
update the memory using the determined update to the memory; and
generate an output for the current time step using the updated memory.
13. The one or more non-transitory computer storage media of claim 12 , wherein the input received at each of the plurality of time step is one of (i) a network input of the memory-based neural network for the current time step or (ii) an intermediate input derived from a network input of the memory-based neural network for the current time step.
14. The one or more non-transitory computer storage media of claim 12 , wherein the memory-based neural network
comprises a plurality of self-attention layers, each self-attention layer corresponding to a respective subset of the memory vectors of the memory, and a respective subset of the memory vectors of the appended memory;
wherein at each time step, applying the attention mechanism over the memory vectors in the memory and memory vectors in the appended memory comprises:
determining the first update to the memory by using one of the plurality of self-attention layers and its respective parameter set, and
determining proposed updates to the memory by using other self-attention layers of the plurality of self-attention layers, wherein each of the other self-attention layers is configured to determine a respective proposed update to the memory by applying an attention mechanism over the respective subset of memory vectors in the memory and the respective subset of memory vectors in the appended memory and
wherein the memory-based neural network is further configured to, at each time step, combine the first update with the proposed updates to generate the update for the memory.
15. The one or more non-transitory computer storage media of claim 14 , wherein the self-attention layers operate in parallel.
16. The one or more non-transitory computer storage media of claim 14 , wherein combining the first update with the proposed updates comprises column-wise concatenating the first update and the proposed updates to generate the update for the memory.
17. The one or more non-transitory computer storage media of claim 14 ,
wherein at each time step, applying the attention mechanism over the memory vectors in the memory and memory vectors in the appended memory comprises:
generating, by the one of the plurality of self-attention layers, from a subset of the memory vectors in the memory, the query matrix, by:
applying a query linear projection of the subset of the memory vectors in the memory onto a query weight matrix to generate the query matrix,
generating, by the one of the plurality of self-attention layers, from a subset of the memory vectors in the appended memory, the key matrix by:
applying a key linear projection of the subset of memory vectors in the appended memory onto a key weight matrix to generate the key matrix,
generating, by the one of the plurality of self-attention layers, from the subset of the memory vectors in the appended memory, the value matrix by:
applying a value linear projection of the subset of memory vectors in the appended memory onto a value weight matrix to generate the value matrix,
generating, from (i) the query matrix generated from the subset of the memory vectors in the memory and (ii) the key matrix generated from the subset of the memory vectors in the appended memory, a weight matrix by:
multiplying the query matrix and a transpose of the key matrix to determine a first temporary matrix,
dividing each element of the first temporary matrix by a scaling factor to determine a second temporary matrix,
applying a softmax operator on each element of the second temporary matrix to generate the weight matrix, and
determining the first update to the memory based on the weight matrix and the value matrix generated from the subset of the memory vectors in the appended memory by:
multiplying the weight matrix and the respective value matrix to determine the first update to the memory.
18. The one or more non-transitory computer storage media of claim 17 , wherein the scaling factor is a square root of the dimension of key vectors in the respective key matrix.
19. The one or more non-transitory computer storage media of claim 14 , wherein the memory-based neural network further comprises a feedforward neural network and one or more memory gating layers; and
wherein at each time step, generating the output for the current time step using the updated memory comprises:
processing the updated memory using the feedforward neural network to generate an intermediate output; and
gating the intermediate output using the one or more memory gating layers to generate the output for the current time step.
20. The one or more non-transitory computer storage media of claim 19 , wherein the feedforward neural network is a multilayer perceptron (MLP).