Patent Description:
Some neural networks are recurrent neural networks. A recurrent neural network is a neural network that receives an input sequence and generates an output sequence from the input sequence. In particular, a recurrent neural network uses some or all of the internal state of the network after processing a previous input in the input sequence in generating an output from the current input in the input sequence.

<CIT> describes a neural network system augmented with external memory, in particular a matrix. Reading data from the external memory involves reading a vector that is a weighted sum of rows of the matrix. <NPL> describes a neural network architecture which uses multi-headed self-attention.

The the system explicitly allows memories, i.e. stored memory vectors, to interact with one another and the input to provide a relational memory core for use, e.g. for relational reasoning in a wide range of tasks. More specifically, the attention mechanism is applied over the plurality of memory locations at a single time step, that is, e.g. not across all previous representations computed from all previous inputs (e.g., observations).

The subject matter described in this specification realizes one or more of the following advantages. By incorporating a new memory module, i.e., a relational memory core (RMC), into a memory-based neural network, the memory-based neural network described in this specification can have the ability to perform complex relational reasoning with information seen at different points in time.

In particular, while conventional memory-based neural networks can model temporal data by leveraging an ability to remember information for long periods, they may struggle at tasks that involve an understanding of the ways in which entities are connected, i.e., tasks that require relational reasoning (hereafter referred to as "relational reasoning tasks") due to the lack of ability to perform complex relational reasoning with the information they remember.

For example, a language modeling task that aims to predict a word given a sequence of observed words is one of the relational reasoning tasks that conventional neural networks would struggle on, as this task requires an understanding of how words that were observed in previous time steps are connected or related to each other. The memory-based neural network described herein can perform such language modeling tasks and therefore are useful for real-world applications such as predictive keyboard and search-phase completion, or can be used as components within larger systems (e.g., machine translation, speech recognition, and information retrieval systems).

As another example, the described memory-based neural network can obtain high performance on reinforcement learning tasks such as controlling an agent interacting with an environment that receives as input data characterizing the environment (observations) and in response to each observation generates an output that defines an action to be performed by the agent in order to complete a specified task. The specified task can be, for example, navigating an environment to collect pre-specified items while avoiding moving obstacles. Such tasks require relational reasoning capability, as the neural network must predict the dynamics of the moving obstacles in memory based on previous observations, and plan the agent's navigation accordingly and also based on remembered information about which items have already been picked up.

Other examples of tasks that require relational reasoning among temporal information in time-series data include multi-lingual machine translation, autonomous navigation (e.g., in self-driving cars), and health monitoring and treatment recommendations.

To address the deficiencies of existing memory-based neural networks, the memory-based neural network described in this specification uses an RMC module that employs multi-head dot product attention to allow memories to interact with each other. In contrast to previous memory architectures, the described memory-based neural network applies attention between memories at a single time step, and not across all previous representations computed from all previous inputs (e.g., observations). Thus, the described memory-based neural network can have a better understating of how memories are related, especially how a new input is related/correlated with previous inputs. Hence, the described memory-based neural network can better solve relational reasoning tasks than conventional neural networks. As a result, systems that implement the described memory-based neural network can have an increased capacity for handling relational reasoning tasks, resulting in improved accuracy when performing such tasks. Further, as the described memory-based neural network is configured to apply a self-attention mechanism over the plurality of memory vectors of the memory and the input at a single time step, information about entities involved in a relational reasoning task are implicitly parsed at each time step. Thus, the systems are more data efficient (compared to systems that employ conventional neural network memory architectures) because they may require less data and fewer learning steps to learn the relational reasoning task. The systems may also generalize better e.g. across multiple different but related tasks.

<FIG> shows an architecture of an example neural network system <NUM> that includes a memory-based neural network <NUM> including a relational memory core (RMC) <NUM>. The neural network system <NUM> is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

Generally, the neural network system <NUM> is a machine learning system that is configured to receive a sequence of inputs and generate a sequence of outputs from the inputs in order to perform a relational reasoning task. A relational reasoning task requires an understanding of the relationship among inputs that were received at previous time steps and the new input received at the current time step and using this understanding to accomplish a higher order goal. As a simplified example, the neural network system <NUM> can receive inputs that are a sequence of randomly sampled vectors. The output of the system at each time step can be an answer to the question: "What is the nth farthest vector in Euclidean distance from vector m?", where the vector values n and m are randomly sampled per sequence.

As another example, the neural network system <NUM> may be a neural machine translation system. In this example, the input sequence is a sequence of words in an original language, e.g., a sentence or phrase, and the output sequence may be a translation of the input sequence into a target language, i.e., a sequence of words in the target language that represents the sequence of words in the original language. Relational reasoning is important for machine translation, especially when a whole sentence or text needs to be translated. In these cases, in order to produce an accurate translation, the system <NUM> needs to understand the meaning of a whole sentence instead of single words, and therefore needs to understand how the words in the input sequence are connected/related to each other.

As another example, the neural network system <NUM> may be a natural language processing or generation system. For example, if the input sequence is a sequence of words in an original language, e.g., a sentence or phrase, the output sequence may be a summary of the input sequence in the original language, i.e., a sequence that has fewer words than the input sequence but that retains the essential meaning of the input sequence. As another example, if the input sequence is a sequence of words that form a question, the output sequence can be a sequence of words that form an answer to the question. In these examples, relational reasoning is required as the system <NUM> needs to understand the meaning of each word in the input sequence and compare the words to each other to determine keywords and the essential meaning of the input sequence. In another example, the input comprises a sequence of words and the output comprises time series data for generating speech corresponding to the words.

As another example, the neural network system <NUM> may be part of a computer-assisted medical diagnosis system. For example, the input sequence can be a sequence of data from an electronic medical record of a patient and the output sequence can be a sequence of predicted treatments. To generate a predicted treatment, the system <NUM> needs to analyze multiple pieces of data in the input sequence to find a correlation between these pieces of data. Based on the correlation, the system <NUM> can determine, for example, symptoms of a disease and/or progression of an existing disease in order to predict an appropriate treatment for the patient. In another example the input may comprise a sequence of sensor data from a medical sensor sensing a condition or one or more parameters of a patient and the output may comprise time series data representing a condition or degree of concern or alert for the patient.

As yet another example, the neural network system <NUM> may be, or be part of, a reinforcement learning system for controlling an agent interacting with an environment that receives as input data characterizing the environment (observations) and in response to each observation, generates an output that defines an action to be performed by the agent in order to complete a specified task. The specified task can be, for example, navigating an environment to collect pre-specified items while avoiding moving obstacles. Such tasks require relational reasoning capability, as the neural network must predict the dynamics of the moving obstacles in memory based on previous observations, and plan the agent's navigation accordingly and also based on remembered information about which items have already been picked up. Some further examples of such a reinforcement learning system are given later.

As another example, the neural network system <NUM> may be part of an image processing or generation system. For example, the input sequence can be an image, i.e., a sequence of color values from the image, and the output can be a sequence of text that describes the image. As another example, the input sequence can be a sequence of text describing a context and the output sequence can be an image that describes the context. These examples also require relational reasoning among components (e.g., color values from the image) in the input sequence in order to produce an accurate output (e.g., text that describes the image). In another example the neural network system <NUM> may be used to implement a recurrent neural network for image generation such as DRAW (arXiv: <NUM>), where the memory may be used for the read and write operations instead of the described selective attention mechanism. In this case the input and output sequences may define images during training, and afterwards the output sequence may define a generated image optionally dependent upon a conditioning variable.

In order to perform the relational reasoning task, the neural network system <NUM> includes a memory <NUM> that is configured to store a respective memory vector at each of a plurality of memory locations of the memory <NUM>. The memory vectors may include information about inputs and outputs of previous time steps and optionally, information about intermediate representations of the inputs that are obtained during the processing of the previous inputs. The memory <NUM> can be denoted as a memory matrix M that has row-wise memory vectors mi. If a current time step is the first time step of a plurality of time steps, the memory matrix M can be randomly initialized or initialized with all of the row-wise memory vectors set to zero. The memory <NUM> can be implemented as one or more physical storage devices in one or more physical locations or as one or more logical storage devices.

To allow for interactions between memory vectors, the memory-based neural network <NUM> includes the RMC <NUM> that is configured to apply an attention mechanism over the memory vectors and an input received at each time step of a plurality of time steps. An attention mechanism generally maps a query and a set of key-value pairs to an output, where the query, keys, and values are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.

At each time step, the RMC <NUM> is configured to receive an input <NUM> and the memory <NUM>. The input <NUM> may be a network input of the memory-based neural network <NUM> for the current time step, or an intermediate input derived from a network input of the memory-based neural network for the current time step. For example, the memory-based neural network <NUM> can directly provide the network input for the time step to the RMC <NUM> or can first process the network input through one or more neural network layers and provide the output of this processing as the input to the RMC <NUM> at the time step.

The RMC <NUM> determines an update to the memory <NUM> by applying a multi-head dot product attention (also known as self-attention) over the memory vectors in the memory <NUM> and the received input <NUM>. By using multi-head dot product attention, each memory vector mi in the memory matrix M can attend over all of the other memory vectors and update its content based on the attended information. Further, because the RMC <NUM> receives both the memory <NUM> and the input <NUM>, the neural network system <NUM> described herein can better understand the relationship between the current input at a current time step and inputs/outputs of previous time steps that are stored in the memory <NUM>.

To apply a multi-head dot product attention, the RMC <NUM> includes multiple self-attention layers <NUM> (i.e., multiple heads). For example, the RMC <NUM> includes h self-attention layers where h ≥ <NUM>. The RMC maintains h sets of parameters for the h self-attention layers. Each of the h self-attention layers corresponds to a respective parameter set of the h parameter sets, a respective sub-memory of the memory matrix M, and a respective sub-memory of an appended memory matrix [M; x] that is created by appending the input <NUM> (denoted as x) to the memory. For example, the input <NUM> is an input vector x, and [M; x] is a row-wise concatenation of the memory matrix M and the input x (i.e., input <NUM>). A sub-memory of the memory matrix M is a subset of the memory vectors in the memory M and there is generally no overlap between the sub-memories of M. Similarly, a sub-memory of the appended memory matrix [M; x] is a subset of the memory vectors in the memory M with the input x appended to the subset. Each set of the h parameter sets, which corresponds to a self-attention layer of the h self-attention layers, includes a respective query weight matrix Wq, a respective key weight matrix Wk, and a respective value weight matrix Wv.

At each time step, each self-attention layer is configured to determine a proposed update to the respective sub-memory by applying an attention mechanism A( ) over memory vectors in the respective sub-memory and the respective sub-appended memory using the respective parameter set. The self-attention layers operate in parallel.

For example, given that M is an N x F dimensional memory matrix and the RMC <NUM> includes two self-attention layers (i.e. two heads), at each time step, the first self-attention layer computes a first proposed update <MAT> and the second self-attention layer computes a second proposed update <MAT>. M<NUM> and M<NUM> are N x F/<NUM> dimensional memory matrices, and θ and φ denote unique parameters for linear projections to produce the queries (Q = MWq), keys (K = MWk) and values (V = MWv) for each memory vectors (i.e., row mi) in the corresponding sub-memories of the memory matrix M.

The process for determining a proposed update to the respective sub-memory by applying an attention mechanism A( ) over memory vectors in the respective sub-memory and the respective sub-appended memory using the respective parameter set are described in more detail below with reference to <FIG>.

After the proposed updates are generated by the self-attention layers, the RMC <NUM> is configured to, at each time step, combine the proposed updates generated by the h self-attention layers to generate the update for the memory <NUM>. In the above example, the RMC <NUM> combines <MAT> and <MAT> to generate the first updated memory <MAT>, where [:] denotes column-wise concatenation. The first updated memory M̃ is a new memory where information is blended across memory vectors based on the weights in their attention weight matrices. A residual is computed based on the first updated memory M̃ and the memory M using a residual connection <NUM>.

The RMC <NUM> further includes a feedforward neural network <NUM> and one or more memory gating layers, e.g. the gating layer <NUM>. The feedforward neural network <NUM> is configured to receive as input the residual computed based on the first updated memory M̃ and the memory M and to process the residual to generate a second updated memory M'. For example, the feedforward neural network <NUM> includes a multilayer perceptron (MLP) and the MLP is applied row-wise to M̃ to generate M'. A second residual connection <NUM> is applied on the second updated memory M' and the input of the feedforward neural network <NUM> (i.e., the residual computed based on M̃ and M) to generate a second residual. The second residual and the memory M are fed as input to the gating layer <NUM>, and are then gated by the gating layer <NUM> and outputted as an output <NUM> and/or a next memory <NUM>. In some cases, the output <NUM> can be provided to the memory-based neural network <NUM> as a network input of the network <NUM> for the current time step. In some other cases, the memory-based neural network <NUM> can process the output <NUM> through one or more neural network layers to generate a network output for the current time step.

The next memory <NUM> can be fed as input to the RMC <NUM> at the next step or to another network component of the memory-based neural network <NUM>.

In particular, the gating layer <NUM> is configured to process the second residual and the memory M to generate, for each memory vector of the memory M, a respective set of gating parameters including a forget gate f, input gate i, and output gate o, in which each gate has a respective weights and bias, by using a linear projection. For each memory vector, the respective input gate i determines how much the memory vector incorporates information from the current input, the respective forget gate f determines how much the memory vector forgets its previous values from the last time step, and the output gate o determines how much the memory vector influences the output of the RMC <NUM> at the current time step. The gating layer <NUM> uses the sets of gating parameters and their weights and bias to generate updated memory vectors in the output <NUM> and/or the next memory <NUM>. The gating layer <NUM> may include a Long Short-Term Memory (LSTM) neural network layer that processes the second residual and the memory M in accordance with a current internal state of the LSTM layer to update the current internal state and to generate the output <NUM> at each time step. The gating layers <NUM> may optionally apply layer normalization on the output of the LSTM layer to generate the output <NUM> at each time step. Layer normalization is described in detail in <NPL>.

In some implementations, the output <NUM> is an output for the relational reasoning task that the neural network system <NUM> is configured to perform. In some other implementations, the output <NUM> is an intermediate output that characterizes a correlation or a relationship between the input <NUM> and information encoded in the memory <NUM> and that is fed to another network component of the memory-based neural network <NUM> for further processing.

Because all of the operations performed by the RMC <NUM> are differentiable, the memory-based neural network <NUM> can be trained using a conventional neural network training technique that is appropriate for the task that the memory-based neural network <NUM> is performing, e.g., an appropriate supervised learning or reinforcement learning training technique. During the training, gradients can be backpropagated to adjust the values of the parameters of the various components of the RMC <NUM>, e.g., of the self-attention layers, the feedforward neural network layers, and the gating layers.

<FIG> is a flow diagram of an example process <NUM> for each self-attention layer to determine a proposed update to a respective sub-memory of a memory by applying an attention mechanism over memory vectors in the respective sub-memory and a respective sub-appended memory using the respective parameter set. The process <NUM> is performed by each self-attention layer at each time step.

While the following equations show that an attention mechanism is applied on memory matrix M and [M; x] for the sake of simplicity, it should be understood that each self-attention layer only applies the attention over memory vectors that are within their respective sub-memory and sub-appended memory in order to generate a respective proposed update <MAT>. The proposed updates generated by the self-attention layers are then column-wise concatenated by a relational memory core to generate the first updated memory M̃.

Each self-attention layer applies a query linear projection of memory vectors mi in the respective sub-memory onto the respective query weight matrix Wq to generate a respective query matrix: MWq (step <NUM>).

Each self-attention layer applies a key linear projection of memory vectors in the respective sub-appended memory onto the key weight matrix Wk to generate a respective key matrix: [M; x] Wk (step <NUM>).

Each self-attention layer applies a value linear projection of the memory vectors in the respective sub-appended memory onto the value weight matrix to generate a respective value matrix: [M; x]Wv (step <NUM>).

Each self-attention layer multiplies the respective query matrix and a transpose of the respective key matrix to determine a first temporary matrix: MWq([M; x]W)T (step <NUM>).

Each self-attention layer divides each element of the first temporary matrix by a scaling factor to determine a second temporary matrix (step <NUM>). The scaling factor is a square root of the dimension dk of key vectors in the respective key matrix.

The second temporary matrix is computed using the following equation: <MAT>.

Each self-attention layer applies a softmax operator on each element of the second temporary matrix to generate a softmax weight matrix: <MAT> (step <NUM>).

Each self-attention layer multiplies the softmax weight matrix and the respective value matrix to determine the proposed update <MAT> for the respective sub-memory (step <NUM>): <NUM> <MAT>.

The proposed updates <MAT> generated by the self-attention layers are then column-wise concatenated by a relational memory core (e.g., the RMC <NUM> of <FIG>) to generate the first updated memory M̃.

<FIG> illustrate an example process for each self-attention layer to determine a proposed update <MAT> to a respective sub-memory of a memory M.

As shown in <FIG>, a linear projection is used to construct queries (Q = MWq), keys (K = [M; x]Wk), and values (V = [M; x]Wv) for each memory vector in the sub-memory of memory M (<NUM>). Next, as shown in <FIG>, the queries Q is used to perform a scale dot-product attention over the keys, K (<NUM>). The returned scalars QKT are put through a softmax-function to generate a set of weights, which are then used to return a weighted average of values from V as: <MAT> where dk is the dimensionality of the key vectors used as a scaling factor. Equivalently: <MAT>.

<FIG> is a flow diagram of an example process for processing an input to generate an output for a current time step using a relational memory core. For convenience, the process <NUM> will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system <NUM> of <FIG>, appropriately programmed in accordance with this specification, can perform the process <NUM>.

The system receives an input (step <NUM>). The input may be a network input of a memory-based neural network (that include the relational memory core) of the system for the current time step, or an intermediate input derived from a network input of the memory-based neural network for the current time step. The system includes a memory M configured to store a respective memory vector at each of a plurality of memory locations in the memory.

The system determines an update to the memory M using the relational memory core (step <NUM>). In particular, the system applies, using the relational memory core, an attention mechanism over the memory vectors in the memory M and the received input to determine an update to the memory M. The attention mechanism is a multi-head dot product attention (or self-attention) which is described in detail above with reference to <FIG> and <FIG>.

The system updates the memory using the determined update to the memory to generate an updated memory M̃ (step <NUM>).

The system generates an output for the current time step using the updated memory (step <NUM>).

In particular, the system applies a residual connection on the first updated memory M̃ and the memory M to generate a first residual. The system processes the first residual using a feedforward neural network to generate a second updated memory M' (step <NUM>). For example, the feedforward neural network <NUM> includes a multilayer perceptron (MLP) and the MLP is applied row-wise to the first residual to generate the second updated memory M'.

The system applies a second residual connection on the second updated memory M' and the input of the feedforward neural network (i.e., the residual computed based on M̃ and M) to generate a second residual.

The system applies a gating operation on the second residual and the memory M using a gating neural network layer to generate a gating output and/or a next memory. In some implementations, the gating output of the gating layer is an output of the system for the time steps. In some other implementations, the system further processes the gating output through one or more neural network layers to generate an output for the current time step. The next memory can be fed as input to the relational memory core at the next time step or to another network component of the system.

As previously described the neural network system may be incorporated into a reinforcement learning system controlling an agent in an environment. Such a reinforcement learning system may be of any type. For example it may be a policy-based system (such as Advantage Actor Critic (A3C), Mnih et al. <NUM>), which directly parameterizes a policy and optionally a value function, or a Q-learning system in which the output approximates an action-value function and optionally a value of a state for determining an action, and it may be a distributed reinforcement learning system such as IMPALA (Importance-Weighted Actor-Learner), <NPL>.

The environment may be a real-world environment and the agent a mechanical/electromechanical agent such as a robot or other static or moving machine interacting with the environment to accomplish a task, e.g. to locate an object of in the environment or to move an object to a specified location in the environment or to navigate to a specified destination in the environment; or the agent may be an autonomous or semi-autonomous land or air or sea vehicle navigating through the environment. The input data (observations) may then include, e.g., one or more of: still or moving images e.g. from a camera or LIDAR, and position, linear or angular velocity, force, torque or acceleration, and global or relative pose of one or more parts of an object and/or the agent. The observations may be defined in <NUM>, <NUM> or <NUM> dimensions, and may be absolute and/or relative, and egocentric or otherwise. The output defining actions may be control inputs to control the robot, e.g., torques for the joints of the robot or higher-level control commands; or to control the autonomous or semi-autonomous land or air or sea vehicle, e.g., torques to the control surface or other control elements of the vehicle or higher-level control commands; or e.g. actions to control steering, acceleration or braking. In another example environment may be a simulated environment corresponding to the real-world environment, and the agent a simulated agent, for training the reinforcement learning system to control the agent e.g. before use in the real-world.

In the case of an electronic agent the observations may include data from one or more sensors monitoring part of a plant or service facility such as current, voltage, power, temperature and other sensors and/or electronic signals representing the functioning of electronic and/or mechanical items of equipment. In some applications the agent may control actions in a real-world environment including items of equipment, for example in a facility such as: a data center, server farm, or grid mains power or water distribution system, or in a manufacturing plant or service facility. The observations may then relate to operation of the plant or facility. For example additionally or alternatively to those described previously they may include observations of power or water usage by equipment, or observations of power generation or distribution control, or observations of usage of a resource or of waste production. The agent may control actions in the environment to increase efficiency, for example by reducing resource usage, and/or reduce the environmental impact of operations in the environment, for example by reducing waste. For example the agent may control electrical or other power consumption, or water use, in the facility and/or a temperature of the facility and/or items within the facility. The actions may include actions controlling or imposing operating conditions on items of equipment of the plant/facility, and/or actions that result in changes to settings in the operation of the plant/facility e.g. to adjust or turn on/off components of the plant/facility. In other applications the agent manages distribution of tasks across computing resources e.g. on a mobile device and/or in a data center. In these implementations, the actions may include assigning tasks to particular computing resources.

While this specification contains many specific implementation details, these should not be construed as limitations, but rather as descriptions of features that may be specific to particular embodiments of particular inventions.

Claim 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 (<NUM>; <NUM>) configured to store a respective memory vector at each of a plurality of memory locations in the memory;
a neural network (<NUM>) that comprises a memory module (<NUM>) for relational reasoning, wherein the neural network is configured to:
at each of a plurality of time steps:
receive an input (<NUM>);
determine an update to the memory;
update the memory using the determined update to the memory; and
generate an output (<NUM>) for the current time step using the updated memory; wherein
determining the update comprises the memory module applying an attention mechanism over the memory vectors in the memory and the received input;
the memory module maintains a plurality of parameter sets; and
the memory module comprises a plurality of self-attention layers (<NUM>), each self-attention layer corresponding to: i) a respective parameter set, ii) a respective sub-memory of the memory that is a subset of the memory vectors in the memory, and iii) a respective sub-memory of an appended memory that is created by appending the received input to the memory; wherein
at each time step, each self-attention layer is configured to determine a proposed update to the respective sub-memory by applying an attention mechanism over memory vectors in the respective sub-memory and the respective sub-memory of the appended memory using the respective parameter set; wherein
the self-attention layers operate in parallel and the attention mechanism is applied over the plurality of memory locations at a single time step; and wherein
the memory module is further configured to, at each time step, combine the proposed updates generated by the plurality of self-attention layers to generate the update for the memory.