Enhancing attention-based neural networks to selectively focus on past inputs

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for performing a machine learning task on a network input as 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 a 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 to (i) a compressed representation in the compressed memory data of the layer, (ii) a hidden state in the episodic memory data of the layer, and (iii) a corresponding hidden state at each of a plurality of input positions in a particular network input to generate a corresponding activation for each input position in the layer input.

Description

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS

[0002] This application claims priority to U.S. Provisional Application No. 62 / 906,078, filed on September 25, 2019. The disclosure of the prior application is considered part of, and incorporated by reference into, the disclosure of the present application. Technical Field

[0003] This specification relates to neural networks. Background Art

[0004] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict outputs given inputs. Some neural networks include one or more hidden layers in addition to the output layer. The output of each hidden layer is used as the input for the next layer in the network (i.e., the next hidden layer or output layer). Each layer of the network generates an output from the input it receives, based on the current values ​​of its corresponding parameter set. Summary of the Invention

[0005] This specification describes a system implemented as a computer program on one or more computers in one or more locations that implements a neural network configured to perform a machine learning task on a network input to generate a network output. Notably, the neural network is a memory-augmented attention neural network that communicates data with one or more memory devices, the one or more memory devices maintaining (i) short-term granular memory, also known as episodic memory, that stores relatively recent (e.g., in terms of time, sequence, or location) information generated by each of one or more layers of the neural network while performing the task, and (ii) long-term coarse memory, also known as compressed memory, that stores older information generated by each of the one or more layers of the neural network while performing the task. When generating a network output from a network input, the system can utilize information stored at both types of memory.

[0006] Generally speaking, one innovative aspect of the subject matter described in this specification can be embodied in a method for processing a sequence of network inputs using a neural network to generate a network output, wherein each network input has a corresponding input at each of a plurality of input positions, the neural network having a plurality of attention layers, each attention layer respectively applying an attention mechanism to a layer input including a corresponding hidden state at each of the plurality of input positions to generate a corresponding activation for each input position. The method includes, for each specific network input in the sequence: for each attention layer in the neural network: maintaining corresponding episodic memory data, the corresponding episodic memory data including corresponding hidden states processed by the attention layer for a first portion of previous network inputs preceding the specific network input in the sequence; maintaining corresponding compressed memory data, the corresponding compressed memory data including compressed representations of corresponding hidden states processed by the attention layer for a second portion of previous network inputs preceding the first portion of previous network inputs in the sequence; receiving a layer input to be processed by the attention layer during processing of the specific network input using the neural network; and applying an attention mechanism to (i) the compressed representation in the compressed memory data for the layer, (ii) the hidden state in the episodic memory data for the layer, and (iii) the corresponding hidden state at each of the multiple input positions in the specific network input to generate corresponding activations for each input position in the layer input.

[0007] The method may also include: updating the episode memory data to include a corresponding hidden state for each input position in the specific network input; and updating the compressed memory data to include a compressed representation of the corresponding hidden state of the earliest network input in the first portion of previous network inputs. Updating the compressed memory data to include a compressed representation of the corresponding hidden state of the earliest network input in the first portion of previous network inputs may include: determining a compression ratio; compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network inputs according to the compression ratio to generate a compressed representation; and modifying the compressed memory data to include the generated compressed representation. The method may also include removing the corresponding hidden state of the earliest network input in the first portion of previous network inputs from the episode memory data. Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network inputs according to the compression ratio may include: applying a max pooling function to the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network inputs, with a stride equal to the compression ratio. Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input according to the compression ratio may include: applying an average pooling function to the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input, wherein the stride is equal to the compression ratio. Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input according to the compression ratio may include: sorting the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input in descending order of corresponding activation weights associated with multiple input positions in the earliest network input and the second portion of previous network input; and discarding the corresponding hidden state of the position in the earliest network input and the corresponding hidden state associated with the lowest activation weight processed by the attention layer for the second portion of previous network input. Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input according to the compression ratio may include: determining a kernel size for a 1D convolution function; and applying the 1D convolution function to the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input, wherein the stride is equal to the compression ratio and the kernel size is equal to the determined kernel size.Compressing the hidden state corresponding to the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network inputs according to the compression rate may include: applying a multi-layer dilated 1D convolution function to the hidden state corresponding to the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network inputs. The specific network input is an input that is preceded by at least a predetermined threshold number of network inputs in the sequence. The method may further include, for each earlier network input in a plurality of earlier network inputs that is preceded by less than the predetermined threshold number of network inputs: for each attention layer in the neural network: maintaining corresponding episode memory data, the corresponding episode memory data including a corresponding hidden state to be processed by the attention layer for the previous network input preceding the earlier network input; receiving a layer input to be processed by the attention layer during processing of the earlier network input using the neural network; applying an attention mechanism to (i) the hidden state in the episode memory data for the layer and (ii) the corresponding hidden state at each of a plurality of input positions in the earlier network input to generate a corresponding activation for each input position in the layer input; and updating the episode memory data to include the corresponding hidden state for each input position in the earlier network input. In some implementations, the predetermined threshold number of network inputs is defined by a predetermined threshold value for the size of the episode memory data. In some implementations, the episodic memory data and the corresponding hidden state for each input position in the particular network input are represented as respective multidimensional arrays; and updating the episodic memory data to include the corresponding hidden state for each input position in the particular network input comprises: concatenating the two multidimensional arrays along a same dimension of the multidimensional array. In some implementations, the neural network may further comprise one or more fully connected layers, one or more layer normalization layers, one or more activation layers, or one or more convolutional layers. The method may further comprise: training the neural network on training data to repeatedly update current values ​​of network parameters, wherein, during training: gradually increasing the number of training neural network inputs between each two consecutive updates. The method may further comprise: applying a stopping gradient during training to (i) the compressed representation in the compressed memory data for the layer and (ii) the hidden state in the episodic memory data for the layer.

[0008] Other embodiments of this aspect include corresponding computer systems, devices, and computer programs recorded on one or more computer storage devices, each of which is configured to perform the actions of the method. The one or more computer systems can be configured to perform specific operations or actions by means of software, firmware, hardware, or any combination thereof installed on the system, which software, firmware, hardware, or any combination thereof can cause the system to perform the actions during operation. The one or more computer programs can be configured to perform specific operations or actions by including instructions that, when executed by a data processing device, cause the device to perform the actions.

[0009] Particular embodiments of the subject matter described in this specification can be implemented so as to realize one or more of the following advantages.

[0010] In general, it can be difficult to train a neural network to capture long-term dependencies within a sequence of network inputs. Conventional approaches to addressing this limitation include maintaining a memory at each attention layer of the neural network that includes a representation of the corresponding hidden state for each input position of the previous network input, thereby enabling the attention layer to focus on longer sequences of network inputs. However, maintaining (i.e., storing and updating) such memory data has various problems. The first is that it consumes a large amount of computational resources (e.g., memory, computing power, or both). The second is that once the maximum size of the memory is reached, the oldest representations in the memory must be discarded, thereby limiting the amount of representations that the network can focus on.

[0011] However, the techniques described in this specification allow the system to maintain compressed memory data at each attention layer of a neural network, the compressed memory data including a compressed representation of the corresponding hidden state for each input position of the oldest network input. In other words, the techniques allow the system to effectively compress the corresponding hidden state for each input position of the oldest network input into the compressed memory.

[0012] By compressing the oldest hidden states rather than discarding them, this technique allows the neural network to consider context, even if the context occurred in the distant past relative to the input currently being processed. Because the oldest hidden states are compressed, this additional context is available with minimal computational overhead. The neural network implementation is able to retain salient information from the past.

[0013] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, drawings, and claims. BRIEF DESCRIPTION OF THE DRAWINGS

[0014] Figure 1 An example attention neural network system is shown.

[0015] Figure 2 is a flowchart of an example process for generating attention layer outputs.

[0016] Figure 3 is a flowchart of an example process for updating the compressed memory associated with an attention layer.

[0017] Figure 4 is a diagram of maintaining the memory associated with the attention layer.

[0018] Like reference numbers and designations throughout the various drawings indicate like elements. DETAILED DESCRIPTION

[0019] This specification describes a system implemented as a computer program on one or more computers in one or more locations that performs machine learning tasks on network inputs.

[0020] A machine learning task can be any machine learning task that: (i) operates on a network input that is an input sequence (i.e., a sequence of inputs each having a corresponding input at each of a plurality of input positions), (ii) generates a network output that is an output sequence, or (iii) both.

[0021] Some examples of machine learning tasks that the system can be configured to perform are as follows.

[0022] As an example, a machine learning task can be neural machine translation, where the input to the neural network is a sequence of text in one language, and the output generated by the neural network can be a score for each text segment in a set of text segments in another language, where each score represents an estimated likelihood that the text segment in the other language is an appropriate translation of the input text into the other language. Thus, for example, each input position can be derived from a word in one language, and the network output can include an output sequence providing a translation of the word into the other language, e.g., having output positions corresponding to the input positions, and where the output positions provide data for the word in the other language.

[0023] As another example, the task may be an audio processing task. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each text segment in a set of text segments, each score representing an estimated likelihood that the text segment is a correct transcription of the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may indicate whether a particular word or phrase ("hot word") is spoken in the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may identify the natural language in which the utterance is spoken. Thus, in general, the network input may include audio data used to perform an audio processing task, and the network output may provide the result of the audio processing task, for example, to recognize a word or phrase or to convert audio into text.

[0024] As another example, the task may be a natural language processing or understanding task that operates on a text sequence of a certain natural language, such as an entailment task, a paraphrase task, a text similarity task, a sentiment task, a sentence completion task, a grammar task, etc.

[0025] As another example, the task may be a text-to-speech task, where the input is text in a natural language or features of text in a natural language, and the network output defines a spectrogram or other data including audio defining the text spoken in the natural language.

[0026] As another example, the task can be a health prediction task, where the input is a sequence derived from a patient's electronic health record data, and the output is a prediction related to the patient's future health, such as a predicted treatment that should be prescribed for the patient, the likelihood that the patient will experience an adverse health event, or a predicted diagnosis for the patient.

[0027] As another example, the task can be a text generation task, where the input is a text sequence and the output is another text sequence, such as a completion of the input text sequence, a response to a question posed in the input sequence, or a text sequence about a topic specified by the first text sequence. As another example, the input of a text generation task can be input other than text, such as an image, and the output sequence can be text describing the input.

[0028] As another example, the task may be an image generation task, where the input is a conditioning input and the output is a sequence of intensity values ​​for pixels of an image.

[0029] As another example, the task can be an agent control task, where the input is a sequence of observations or other data representing the state of the environment, such as a video sequence, and the output defines the action that the agent performs in response to the most recent data in the sequence. The agent can be a mechanical agent, such as a real-world or simulated robot, a control system for an industrial facility, or a control system that controls different kinds of agents.

[0030] To perform machine learning tasks, the system includes an attention neural network comprising a plurality of layers, each layer operating on a respective input sequence comprising a respective input vector at each of one or more positions.

[0031] In addition, each of some or all of the layers includes an attention layer, and in some implementations, a feedforward layer. As used herein, an attention layer is a neural network layer that includes an attention mechanism, which can optionally be a self-attention mechanism, a masked attention mechanism, and / or a multi-head attention mechanism. The attention layer receives an input sequence for the layer and applies an attention mechanism to the input sequence for the layer to generate an focused input sequence.

[0032] The exact attention mechanism applied by the attention layer depends on the configuration of the attention neural network, but in general, the attention mechanism maps a query and a set of key-value pairs to an output, where the query, key, and value 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 between the query and the corresponding key. More specifically, each attention layer can apply a scaled dot product attention mechanism. In scaled dot product attention, for a given query, the attention layer computes the dot product of the query with all keys, divides each dot product by a scaling factor, e.g., by the square root of the dimension of the query and key, and then applies a softmax function to the scaled dot products to obtain the weights for the values. The attention layer then computes a weighted sum of the values ​​based on these weights. Thus, for scaled dot product attention, the compatibility function is a dot product, and the output of the compatibility function is further scaled by the scaling factor.

[0033] The use of an attention mechanism allows a neural network to relate different positions of a single sequence in order to compute a representation of the sequence, effectively learning dependencies between distant positions during training. This can improve the accuracy of the system when using neural networks for various machine learning tasks that require sequential network input.

[0034] When a feed-forward layer is included, then the feed-forward layer operates on the input sequence of interest to generate the output sequence of the layer.

[0035] In general, the layers within an attention neural network can be arranged in any of a variety of configurations.

[0036] As an example, when the network input is an input sequence, the attention neural network can include an encoder neural network that includes a subset of multiple layers and encodes the input sequence to generate a corresponding encoded representation of each input in the sequence. In this example, the attention mechanism applied by the layers in the encoder is a self-attention mechanism, such as a multi-head self-attention mechanism. In a self-attention mechanism, the input vector and memory vector operated by the attention mechanism are the same, that is, the vector in the input sequence of the layer.

[0037] As another example, an attention neural network includes a decoder neural network that includes different subsets of multiple layers and processes network inputs or encoded representations of network inputs to generate network outputs.

[0038] In some of these examples, when the network output is an output sequence, the decoder neural network operates autoregressively to generate outputs in the output sequence, and attention sublayers within some or all layers of the decoder apply masked self-attention to the partially generated output sequence. In masked self-attention, the input vector and the memory vector operated by the attention mechanism are the same, but the attention mechanism is masked such that any given position in the input sequence does not attend to any positions after the given position in the input sequence.

[0039] When a neural network includes both an encoder and a decoder, some layers in the decoder apply crisscross attention to the encoded representation, while other layers apply self-attention to the masked or unmasked output sequence. When crisscross attention is applied, the input vector is from the input sequence to the layer, and the memory vector is the encoded representation generated by the encoder.

[0040] When an attention neural network includes a decoder neural network that operates directly on the input sequence (i.e., includes only a decoder and no encoder), the attention layer within the decoder can apply a self-attention mechanism to the input sequence.

[0041] Specific examples of attention-based neural network architectures that include multiple attention layers and that can be modified to include attention layers of the type described in this specification are described in Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171–4186, 2019; and in Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context.In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978–2988, Florence, Italy, July 2019. Association for Computational Linguistics. doi:10.18653 / v1 / P19-1285. URL https: / / www.aclweb.org / anthology / P19-1285; and Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, pp. 5998–6008, 2017. URL https: / / papers.nips.cc / paper / 7181-attention-is-all-you-need.pdf. The entire disclosure of these is incorporated herein by reference in its entirety.

[0042] Figure 1 Illustrated is an example neural network system 100. Neural network system 100 is an example of a system implemented as a computer program on one or more computers in one or more locations in which the systems, components, and techniques described below may be implemented.

[0043] Neural network system 100 may receive input 102 and perform a machine learning task on input 102 to generate output 152 .

[0044] As described above, neural network system 100 can perform any of a variety of tasks involving (i) operating on input 102 as an input sequence, (ii) generating output 152 as an output sequence, or (iii) both, and more particularly, tasks requiring reasoning over long program sequences (e.g., long documents, e.g., multiple consecutive articles or full-length books, or long sequences of observations generated as an agent interacts with an environment).

[0045] The neural network system 100 includes an attention neural network 110, which includes an attention layer 120. The attention layer 120 operates on an input sequence 112 and generates a corresponding output sequence 122.

[0046] In an implementation, input 102 comprises a sequence of network inputs, each network input having a corresponding input at each of a plurality of input positions. Thus, an input sequence 112 can be derived from the network input. In general, a network output 152 depends on an output sequence 122 from one or more attention layers. The network output 152 can provide an output corresponding to each input position. For example, in a natural language processing system, this can be used to determine, for example, a word for a translation of an input, or in a reinforcement learning system to determine an action to be performed at a time step. In a reinforcement learning system, each input position can correspond to a time step at which an observation is received.

[0047] However, in some implementations, network output 152 may provide outputs corresponding to multiple or all input locations. For example, input 102 may include audio or video input, the input locations may be defined by samples of audio or frames of video, and network output 152 may characterize (e.g., classify) information in the audio or video input, such as to identify sounds such as words, or actions or one or more objects depicted by a video.

[0048] Although for convenience, as mentioned above Figure 1 An attention layer is depicted in Fig. 1, but the attention neural network 110 may include other layers, including, for example, an embedding layer, (one or more) output layers, and more attention layers. Other layers that may be included are (non-linear) activation layers, fully connected layers, and layer normalization layers (arXiv:1607:06450).

[0049] In general, the input sequence 112 can be any intermediate sequential data generated by the attention neural network 110 when performing a machine learning task on the input 102. Each input sequence 112 can be a portion of the system input 102 or a fragment of the overall sequence derived from the system input 102. As the system moves through processing different portions of the system input 102 to perform the machine learning task, different input sequences 112 can be derived. For example, the input sequence 112 can be an embedded (i.e., digital) representation of the system input 102 or a fragment of the system input 102 generated by an embedding layer or, more generally, an embedding neural network. Optionally, the embedded representation of the input sequence can be combined with an encoding of the corresponding input position. As another example, the input sequence 112 can be an output sequence generated by a previous attention layer or another layer in the attention neural network 110.

[0050] Specifically, the input sequence 112 has a corresponding hidden state input at each of the plurality of input positions in the input order, and the output sequence 122 has a corresponding activation output at each of the plurality of output positions in the output order. That is, the input sequence 102 has the plurality of hidden state inputs arranged according to the input order, and the output sequence 122 has the plurality of activation outputs arranged according to the output order. Therefore, where the attention neural network includes a stack of multiple attention layers, the hidden state in the input sequence of each attention layer may generally be the output activation generated by a previous attention layer in the stack that precedes the attention layer in the attention neural network 110.

[0051] A common weakness of attention-based neural networks (even those enhanced with memories of reusable information generated when processing previous network inputs in a sequence (e.g., past activation outputs at each network layer)) is the ability to capture long-term dependencies within a sequence of network inputs. In general, applying an attention mechanism at each attention layer of an attention neural network to an input sequence of arbitrarily long length (e.g., an input sequence derived from a system input 102 comprising a long document (e.g., a full-length book) having thousands or millions of characters) can be very expensive and therefore capacity-limited. This is due to the computational cost of attending to each input in an arbitrarily long sequence, and due to the storage cost of retaining this large memory if memory is used.

[0052] Therefore, to help the attention layer 120 process the input sequence 112, the neural network system 100 maintains (e.g., at one or more physical or logical storage devices) a segment memory 130 and a compressed memory 140. In the case where the attention neural network 110 includes multiple attention layers, the neural network system 100 can maintain a single memory, namely, the segment memory 130 or the compressed memory 140, for all attention layers, or maintain different memories for different attention layers. These memories can, but need not, have the same structure, for example, a first-in-first-out (FIFO)-like structure.

[0053] The episodic memory 130 can be implemented as one or more logical or physical storage devices and stores previous layer inputs processed by the attention layer 120 when the system operates on an earlier portion of the system input 102. For example, when the system input 102 is an input sequence, the earlier portion of the system input 102 may include a first portion of the previous system input that precedes the current system input in the sequence, and one or more earlier input sequences of the attention layer 120 that can be derived from the first portion of the previous system input. That is, the episodic memory 130 stores "episodic memory data," e.g., short-term granular data, which includes, at corresponding episodic memory slots, the hidden states 124 of the attention layer 120 from one or more earlier input sequences to immediately preceding the current input sequence 112.

[0054] The compressed memory 140 stores the second portion of the previous layer input, which was processed by the attention layer 120 when the system was operating the earliest portion of the system input 102 that preceded the earlier portion of the system input 102. Following the above example, the earliest portion of the system input 102 may include the second portion of the previous system input as an input sequence that further precedes the first portion of the previous system input in the system input 102, and one or more earliest input sequences of the attention layer 120 may be derived from the second portion of the previous system input. That is, the episode memory 130 stores "compressed memory data" that includes, at corresponding compressed memory slots, hidden states 134 from one or more earliest input sequences to the attention layer 120, the one or more earliest input sequences further preceding one or more earlier input sequences to the attention layer 120 that immediately precede the current input sequence 112.

[0055] As the system 100 moves through processing different portions of the system input 102 to perform machine learning tasks, the system 100 can determine new compressed memory data from hidden states currently stored at the episodic memory 130, hidden states already stored at the compressed memory 140, or both, by utilizing a compression engine 160, where the compression engine 160 can be implemented as, for example, a compression layer of the attention neural network 110, i.e., a network layer configured to apply a compression function to a layer input to output a compressed representation of the layer input.

[0056] The following will refer to Figures 2 to 4 Processing each of the plurality of network inputs in the current input sequence 112 while utilizing the corresponding hidden states 128 and 138 stored at the episodic memory and the compressed memory is described in greater detail.

[0057] In short, to generate a current output sequence 122 from the current input sequence 112 and the hidden states stored in memory, the attention layer 120 is configured to: apply a learned query linear transformation to each hidden state at each input position in the current input sequence to generate a corresponding query Q for each input position; apply a learned key linear transformation to (i) each hidden state at each input position and (ii) each hidden state at each slot in the episodic memory and compressed memory to generate a corresponding key K for each input position and each memory slot; and apply a learned value linear transformation to (i) each hidden state at each input position and (ii) each hidden state at each slot in the episodic memory and compressed memory to generate a corresponding value V for each input position and each memory slot. The attention layer 120 then applies the above-described attention mechanism using these queries, keys, and values ​​to determine an output sequence 122 for the input sequence 104. The output sequence 122 typically includes a corresponding attention vector for each hidden state input at each input position. Generally, the query Q, key K, and value V are all vectors. As used in this specification, the term "learned" means that an operation or value has been adjusted during system training.

[0058] In some implementations, to allow the attention layer to jointly focus on information from different representation subspaces at different locations, the attention layer employs multi-head attention.

[0059] To implement multi-head attention, the attention layer 120 applies h different attention mechanisms in parallel. In other words, the attention layer includes h different attention sub-layers, and each attention sub-layer within the same attention layer receives the same original query Q, original key K, and original value V.

[0060] Each attention sub-layer is configured to transform the original query, key, and value using a learned linear transformation, and then apply the attention mechanism to the transformed query, key, and value. Each attention sub-layer will typically learn different transformations from other attention sub-layers in the same attention layer.

[0061] Specifically, each attention sub-layer is configured to apply the learned query linear transformation to each original query to generate a layer-specific query for each original query, apply the learned key linear transformation to each original key to generate a layer-specific key for each original key, and apply the learned value linear transformation to each original value to generate a layer-specific value for each original value. The attention sub-layer then applies the above-mentioned attention mechanism using these layer-specific queries, keys, and values ​​to generate an initial output for the attention sub-layer.

[0062] The attention layer then combines the initial outputs of the attention sub-layers to generate the final output of the attention layer. Specifically, the attention layer can compute a cascade of the outputs of the attention sub-layers and apply the learned linear transformation to the output of the cascade to generate as output the sequence of attention inputs 124. In some cases, the learned transformation applied by the attention layer reduces the dimensionality of the original keys and values, and optionally the query. For example, when the dimensionality of the original keys, values, and queries is d and there are h attention sub-layers in the attention layer, the attention layer can reduce the dimensionality of the original keys, values, and queries to d / h. This will keep the computational cost of the multi-head attention mechanism similar to the cost of performing the attention mechanism once with full dimensionality, while increasing the representational power of the attention layer.

[0063] Figure 2 is a flow chart of an example process 200 for generating an attention layer output. For convenience, process 200 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., Figure 1 The neural network system 100) can perform process 200.

[0064] The system can repeat process 200 for each sequence of layer inputs. For convenience, each of steps 202-208 will be described as being performed by each attention layer in the attention neural network and for the "current" sequence of layer inputs.

[0065] Generally speaking, a layer input sequence includes a corresponding hidden state input at each of a plurality of input positions. As similarly described above, a layer input sequence can be any intermediate sequential data generated by an attention neural network when performing a machine learning task on a system input to generate a system output. As the system progresses by processing different portions of the system input to perform the machine learning task, different layer input sequences can be derived.

[0066] For example, the layer input sequence can be a representation of the system input embedding generated by the embedding layer. As another example, the layer input sequence can be a sequence of outputs generated by a previous attention layer or other layer in the attention neural network. For example, the corresponding hidden state inputs in the sequence can be multiple activation outputs generated by a previous attention layer in the attention layer stack that immediately precedes the given attention layer in the attention neural network.

[0067] The system maintains corresponding episode memory data (202) for an attention layer in an attention neural network, the corresponding episode memory data comprising corresponding hidden states processed by the attention layer when the system operated on an earlier portion of the system input.

[0068] For example, the earlier portion of the system input may include the first portion of the previous system input that precedes the current portion of the system input in the sequence. For example, the system may do so by maintaining an episode memory of a fixed size that stores a corresponding hidden state at each of a plurality of input positions in one or more previous layer input sequences that have been processed by the attention layer when operating on the first portion of the system input (i.e., before processing the current layer input sequence).

[0069] The system maintains corresponding compressed memory data (204) for an attention layer in an attention neural network, the corresponding compressed memory data comprising compressed representations of corresponding hidden states processed by the attention layer when the system operates on the earliest portion of the system input.

[0070] For example, the earliest portion of the system input may include a second portion of the previous system input that further precedes the first portion of the system input. For example, the system may do so by maintaining a compressed memory that stores a corresponding hidden state at each of a plurality of input positions in one or more earliest layer input sequences that further precede one or more earlier layer input sequences for which the hidden state is stored in the episodic memory.

[0071] The following will refer to Figures 3 and 4 As described above, the system can use any of a variety of techniques to select multiple input positions (hidden states to be removed from the segment memory) from all input positions of one or more earlier layer input sequences. For example, the hidden state in the oldest input sequence within the one or more earlier layer input sequences is always removed, and a compressed representation of the hidden state in the oldest input sequence is stored in the compressed memory. Collectively, the selected hidden states to be stored in the compressed memory are referred to as the compressed representation of the corresponding hidden state processed by the attention layer for the second portion of the previous system input.

[0072] When performing a machine learning task on a current portion of a system input, the system receives, at an attention layer in an attention neural network, a current layer input sequence to be processed by the attention layer to generate a system output using the attention neural network (206). The current layer input sequence may have a corresponding hidden state input at each of a plurality of input positions.

[0073] For episodic or compressed memory, if the current layer input sequence is the first sequence processed by the attention neural network when performing a machine learning task on the network input, the corresponding hidden states maintained at the memory can each have a corresponding predetermined initial value, such as zero, negative or positive infinity, or some other predetermined value. Then, as the system moves through processing new input sequences to perform the machine learning task, these predetermined initial values ​​are gradually replaced with updated values ​​specified by the corresponding hidden states generated by the system.

[0074] The system applies an attention mechanism to (i) a compressed representation of the corresponding hidden state in the compressed memory, (ii) the corresponding hidden state in the episodic memory, and (iii) the corresponding hidden states at multiple input positions in the current layer input sequence to determine a corresponding layer output sequence (208) for the current layer input sequence.

[0075] In more detail, the system can, for example, by using an attention layer or another system component: apply a learned query linear transformation to each hidden state at each input position in the current input sequence to generate a corresponding query Q for each input position; apply a learned key linear transformation to the concatenated representations of (i) each hidden state at each input position and (ii) each hidden state at each slot in the episodic memory and compressed memory to generate a corresponding key K for each input position and each memory slot; and apply a learned value linear transformation to the concatenated representations of (i) each hidden state at each input position and (ii) each hidden state at each slot in the episodic memory and compressed memory to generate a corresponding value V for each input position and each memory slot. In various implementations, each representation can be in the form of a multidimensional array (e.g., a vector), and the concatenated representation can be obtained by concatenating two multidimensional arrays along the same dimension of the multidimensional array. The system can then use these queries, keys, and values ​​to apply the above-described attention mechanism to determine an attention input sequence for the input sequence. The output sequence typically includes a corresponding attention vector for each hidden state at each input position in the current layer input sequence.

[0076] When the attention layer implements multi-head attention, each attention sub-layer in the attention layer is configured to: apply the learned query linear transformation to each layer input at each input position in the current input sequence to generate a corresponding query Q for each input position; apply the learned key linear transformation to the concatenated representation of (i) each hidden state at each input position and (ii) each hidden state at each slot in the episodic memory and the compressed memory to generate a corresponding key K for each input position and each memory slot; apply the learned value linear transformation to the concatenated representation of (i) each hidden state at each input position and (ii) each hidden state at each slot in the episodic memory and the compressed memory to generate a corresponding value V for each input position and each memory slot; then apply the attention mechanism (i.e., the scaled dot product attention mechanism described above) using the query, key, and value to determine the initial attention sub-layer output for each input position and each memory slot. The attention layer then combines the initial outputs of the attention sub-layers as described above.

[0077] In an implementation where each attention layer in turn includes a feed-forward layer, the system can use the feed-forward layer to operate on the focused input sequence to generate an output sequence for the attention layer. Alternatively, the system can provide the focused input sequence as the output sequence for the attention layer.

[0078] The system then proceeds to update the episode memory and, optionally, the compressed memory based on the layer input sequence. Specifically, the system may first update the episode memory to include the corresponding hidden state in the current layer input sequence, and then determine whether the episode memory is "full," i.e., whether all available memory slots within the episode memory have been updated with corresponding values ​​of the hidden state in the current layer input sequence as a result of performing process 200 one or more times.

[0079] In response to a negative determination, i.e., if there are fewer than a predetermined threshold number of layer input sequences preceding the current layer input sequence, the system can continue to update the episode memory data. In some implementations, the predetermined threshold number is defined by a predetermined threshold value for the size of the episode memory (e.g., in terms of available memory slots, each memory slot being operable to store a corresponding hidden state value).

[0080] Specifically, the system can do so by updating the corresponding episode memory slots to include the values ​​of the hidden states that have been processed by the attention layer in the current layer input sequence to generate the corresponding layer output sequence.

[0081] Alternatively, in response to a positive determination, i.e., where there are more than a predetermined threshold number of layer input sequences preceding the current layer input sequence, the system may proceed to process 300 to update the compressed memory data, i.e., in addition to updating the segment memory after executing process 200.

[0082] Figure 3 is a flow chart of an example process for updating a compressed memory associated with an attention layer. For convenience, process 300 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., Figure 1 The neural network system 100) can perform process 300.

[0083] The system determines a compression ratio c (302), which generally defines a measure of the relative reduction in size of the data representation resulting from the compression operation. Higher values ​​of the compression ratio generally result in a greater reduction in the size (i.e., in terms of the number of hidden states) of the selected portion of the episodic memory data to be compressed. For example, the system may receive input specifying a value for the compression ratio from a system user, such as through an application programming interface (API) available to the system. As another example, the system may select a value for the compression ratio, such as based on a predefined compression scheme.

[0084] The system performs compression operation f according to the compression ratio c c , to compress corresponding hidden states in one or more earlier layer input sequences processed by the attention layer when operating on the first portion of the system input to generate a compressed representation (304). Specifically, the system can determine the exact size of the subset of the one or more earlier layer input sequences (i.e., how many hidden states to remove from the segment memory) based on the compression operation, the segment memory, or the size of the compression memory, or a combination thereof. Once determined, the system can use any of a variety of techniques to perform the compression operation, update the compression memory to include the compressed representation, and then remove the corresponding hidden states from the segment memory.

[0085] In some implementations, the system can apply a max pooling function to the corresponding hidden states processed by the attention layer for the subset of the one or more earlier layer input sequences, where the stride is equal to the compression rate. That is, the compression operation can be a max pooling operation, and the compressed representation is the output of the max pooling function computed using the hidden states stored at the episode memory.

[0086] In some implementations, the system can apply an average pooling function to the corresponding hidden states processed by the attention layer for the subset of the one or more earlier layer input sequences, where the stride is equal to the compression rate.

[0087] In some implementations, the system may sort the corresponding hidden states processed by the attention layer for a subset of the one or more earlier layer input sequences in descending order of the corresponding values ​​of the hidden states associated with multiple input positions in the one or more earlier layer input sequences, and thereafter discard the corresponding hidden states of the positions in the subset of the one or more earlier layer input sequences that are associated with the lowest hidden state values.

[0088] In some implementations, the system can determine a kernel size for a 1-D convolution function and then apply the 1-D convolution function to corresponding hidden states processed by the attention layer for a subset of the one or more earlier layer input sequences, where the stride is equal to the compression ratio and the kernel size is equal to the determined kernel size.

[0089] In some implementations, the system may instead apply the derivative of a traditional convolution function (e.g., a dilated 1-D convolution function) to the corresponding hidden states processed by the attention layer for a subset of the one or more earlier layer input sequences.

[0090] In some of these implementations, the system may instead compress, i.e., by applying a compression operation to both (i) the corresponding hidden states in the one or more earlier layer input sequences and (ii) the corresponding hidden states in the one or more earliest layer input sequences that have been processed by the attention layer, or more specifically, the corresponding hidden states processed by the attention layer when the system operates on a second portion of the system input (further prior to the first portion of the system input). In such an implementation, the system may determine an overall compressed representation of the hidden states and the previously compressed hidden states stored at the episode memory and the compression memory, respectively.

[0091] The system modifies the compressed memory data to include a compressed representation (306) by replacing the corresponding current value stored at the compressed memory slot with the updated hidden state value specified by the compressed representation generated from step 304. The corresponding hidden state based on which the compressed representation was generated is then discarded from the segment memory to make room for new hidden states in the upcoming layer input sequence. As new compressed hidden states become available during operation of the system, previously compressed hidden states can be discarded from the compressed memory, e.g., with the oldest hidden states discarded first. However, in an implementation, some or all of the compressed hidden states are retained in the compressed memory throughout the execution of the machine learning task.

[0092] Figure 4 is a diagram of maintaining the memory associated with the attention layer. Figure 4 In the example, the attention neural network includes three attention layers, each of which is configured to receive a length n s=3, i.e., an input sequence 410 (“current sequence”) with corresponding hidden states at each of the three input positions in the input order. Each attention layer is associated with size n m = 6 episode memories 420 (ie, memories configured to store corresponding episode hidden states at each of the six memory slots) and size n cm =6 compression memory 430 (ie, a memory configured to store a corresponding compressed hidden state at each of the six memory slots). Figure 4 In the example of , the compression ratio c = 3, that is, every three episode hidden states are compressed into a single compressed hidden state.

[0093] For each attention layer, the set of three hidden states generated from processing the current sequence is moved to the corresponding memory slots in the episode memory 420. Then, according to the layer-specific compression function f c The episode hidden state (enclosed by a rectangular box) already maintained at episode memory 420 is compressed and moved to a single memory slot within compressed memory 430 .

[0094] An example algorithm for maintaining the memory associated with an attention layer is shown below.

[0095]

[0096]

[0097] In the example algorithm shown above, n m and n cm is the number of corresponding memory and compressed memory slots in each attention layer of the attention neural network. The entire input sequence S = x1, x2, ..., x |s| represents the input to the system (e.g., tokens from a book). These are divided into blocks of size n s The attention neural network receives the fixed-size window of t at time t. (exist Figure 4 When the attention neural network moves to the next sequence, its n s hidden states are pushed into a fixed-size first-in-first-out (FIFO) memory (in Figure 4 The oldest n in memory is called "episode memory" 420). s hidden states are removed, using the compression operation ∫ c : Processing, n s The oldest memory is mapped to the auxiliary FIFO memory to be stored (in Figure 4 is called "compressed memory" 430) compressed memory. d represents the hidden size of the hidden state. c refers to the compression ratio, where higher values ​​indicate coarser-grained compressed memory.

[0098] Process 200 or 300 can be performed on each layer input sequence to generate a layer output sequence for the layer input sequence, from which the system output can be derived. The sequence of layer inputs can be a sequence for which the expected output (i.e., the output sequence that should be generated by the attention layer for the input sequence) is unknown. The system can also perform process 200 or 300 on inputs in a training dataset (i.e., a set of inputs for which the output that should be predicted by the attention layer is known) in order to train the system, i.e., determine training values ​​for the parameters of the attention neural network and, in some implementations, determine any additional parameters required to maintain one or more memories associated with the network.

[0099] During training of the attention neural network, process 200 or 300 can be repeatedly performed on inputs selected from the training data set as part of a conventional machine learning training technique to train the attention neural network, such as Adam or Adam with backpropagation through time, which generally involves iteratively updating corresponding parameter values ​​of the attention neural network based on the calculated gradient of the objective function.

[0100] In some implementations, in addition to one or more terms that penalize the system for generating incorrect system outputs, the objective function being optimized also includes one or more terms that penalize the system for information loss caused by compression. In this way, the attention neural network can be trained to generate high-quality system outputs by efficiently reasoning within the system input or on corresponding inputs derived from the system input. For example, the objective function includes one or more terms that evaluate auxiliary compression loss. For example, such auxiliary compression loss can be a lossless compression target measured in terms of the difference between the reconstruction of the compressed memory content and the original uncompressed episodic memory content. As another example, such auxiliary compression loss can be a lossy compression target measured in terms of the difference between content-based attention (i.e., the attention layer output generated from paying attention to the episodic memory content) and memory-based attention (i.e., the attention layer output generated from paying attention to the compressed memory content).

[0101] Specifically, an objective function training system that uses a penalty system for incorrect system outputs and an objective function training system that uses a penalty system for losing attention information due to a suboptimal compression strategy can be performed jointly, for example, by backpropagating the compression loss gradient into the network parameters of the attention neural network, or separately, for example, by applying a stop gradient operator to the compressed representation maintained in the compression memory of the attention layer, the hidden state maintained in the episodic memory data of the layer, or both. The stop gradient operator prevents the gradient associated with the compression loss from being propagated into the network parameters of the attention neural network, thereby focusing on task-related training of the attention neural network.

[0102] The system can also perform process 200 or 300 in conjunction with any of a variety of other training techniques that are deemed advantageous (e.g., by a system operator) in terms of wall-clock time or computational resources (e.g., memory, computing power, or both). For example, the system can train the system according to a dynamic parameter update frequency, for example, by gradually increasing the number of training neural network inputs to be processed by the attention neural network between each two consecutive parameter value updates. Some implementations of the system can clip gradients during training and / or can use a learning rate schedule that has a warm-up phase that increases during the warm-up phase and then a decay phase.

[0103] An example algorithm for computing an objective function evaluating the attention-reconstruction loss for use in training the system is shown below.

[0104]

[0105] This specification uses the term "configured" in conjunction with system and computer program components. With respect to a system of one or more computers to be configured to perform a particular operation or action, this means that the system has installed thereon software, firmware, hardware, or a combination thereof that, when operated, causes the system to perform the operation or action. With respect to one or more computer programs to be configured to perform a particular operation or action, this means that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0106] Embodiments of the subject matter and functional operations described in this specification may be implemented in digital electronic circuitry, in tangibly embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or a combination of one or more thereof. Embodiments of the subject matter described in this specification may be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium, for execution by a data processing apparatus or for controlling the operation of a data processing apparatus. A computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more thereof. Alternatively or additionally, program instructions may be encoded on an artificially generated propagated signal, such as a machine-generated electrical, optical, or electromagnetic signal, generated to encode information for transmission to a suitable receiver device for execution by the data processing apparatus.

[0107] The term "data processing apparatus" refers to data processing hardware and encompasses all types of apparatus, devices, and machines for processing data, including, for example, a programmable processor, a computer, or multiple processors or computers. The apparatus may also be or include dedicated logic circuitry, such as an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, the apparatus may optionally include code that creates an execution environment for a computer program, for example, code constituting processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of these.

[0108] A computer program, which may also be referred to or described as a program, software, software application, app, module, software module, script, or code, may be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it may be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program may be stored in a portion of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), in a single file dedicated to the program in question, or in multiple coordinated files (e.g., files storing one or more modules, subroutines, or code portions). A computer program may be deployed to execute on one computer or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communications network.

[0109] In this specification, the term "database" is used broadly to refer to any collection of data: the data need not be structured in any particular way, or at all, and it may be stored on a storage device in one or more locations. Thus, for example, an index database may include multiple collections of data, each of which may be organized and accessed differently.

[0110] Similarly, in this specification, the term "engine" is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally speaking, an engine will be implemented as one or more software modules or components installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a specific engine; in other cases, multiple engines can be installed and run on the same one or more computers.

[0111] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, such as an FPGA or ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0112] Computers suitable for executing computer programs include, for example, general or special microprocessors or both, or any other type of central processing unit. Generally speaking, the central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The basic elements of a computer are a central processing unit for executing or running instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by or incorporated into a dedicated logic circuit. Generally speaking, a computer will also include one or more large-capacity storage devices (such as magnetic disks, magneto-optical disks, or optical disks) for storing data, or be operably coupled to receive data from one or more large-capacity devices or transfer data to one or more large-capacity devices or both. However, a computer does not need to have such a device. In addition, a computer can be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, such as a universal serial bus (USB) flash drive, to name a few.

[0113] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, including, for example, semiconductor memory devices, such as EPROM, EEPROM, and flash memory devices; magnetic disks, such as internal hard disks or removable disks; magneto-optical disks; and CD ROM and DVD-ROM disks.

[0114] To provide for interaction with a user, embodiments of the subject matter described in this specification may be implemented on a computer having a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) for displaying information to the user, and a keyboard and pointing device (e.g., a mouse or trackball) for the user to provide input to the computer. Other types of devices may also be used to provide for interaction with the user (e.g., the feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback); and input from the user may be received in any form, including acoustic, voice, or tactile input. Furthermore, a computer may interact with a user by sending documents to and receiving documents from a device used by the user (e.g., by sending a web page to a web browser on a user's device in response to a request received from the web browser). Furthermore, a computer may interact with a user by sending text messages or other forms of messages to a personal device (e.g., a smartphone running a messaging application) and receiving response messages from the user in return.

[0115] The data processing apparatus for implementing a machine learning model may also include, for example, dedicated hardware accelerator units for processing the common and computationally intensive parts of machine learning training or production, i.e., inference, workloads.

[0116] You can use a machine learning framework (for example, the TensorFlow framework, the Microsoft Cognitive Toolkit framework, the Apache Singa framework, or the Apache MXNet framework) to implement and deploy the machine learning model.

[0117] Embodiments of the subject matter described in this specification may be implemented in a computing system that includes a back-end component (e.g., as a data server); or includes a middleware component (e.g., an application server); or includes a front-end component (e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification); or includes any combination of one or more such back-end, middleware, or front-end components. The components of the system may be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include local area networks ("LANs") and wide area networks ("WANs") (e.g., the Internet).

[0118] A computing system may include a client and a server. The client and server are typically remote from each other and typically interact via a communication network. The relationship between the client and the server arises by means of computer programs running on respective computers and having a client-server relationship with each other. In some embodiments, the server transmits data (e.g., an HTML page) to a user device, for example, for the purpose of displaying data to a user interacting with the device acting as a client and receiving user input from the user. Data generated at the user device, for example, the results of a user interaction, may be received at the server from the device.

[0119] Although this specification contains many specific implementation details, these should not be interpreted as limitations on the scope of any invention or the scope that may be claimed, but rather should be interpreted as descriptions of features that may be specific to a particular embodiment of a particular invention. Certain features described in this specification in the context of separate embodiments may also be implemented in combination in a single embodiment. Conversely, the various features described in the context of a single embodiment may also be implemented in multiple embodiments individually or in any suitable sub-combination. In addition, although features may be described above as working in certain combinations and even initially claimed as such, one or more features from the claimed combination may in some cases be deleted from the combination, and the claimed combination may be directed to a sub-combination or a variation of the sub-combination.

[0120] Similarly, although operations are depicted in a particular order in the drawings and recited in the claims, this should not be understood as requiring that the operations be performed in the particular order shown, or in sequential order, or that all illustrated operations be performed, to achieve the desired results. In certain circumstances, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0121] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve the desired results. As an example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method for processing an input comprising a sequence of network inputs using a neural network to generate a network output, each network input in the sequence of network inputs having a corresponding input at each of a plurality of input positions, the neural network having a plurality of attention layers, each attention layer separately applying an attention mechanism to a layer input comprising a corresponding hidden state at each of the plurality of input positions to generate a corresponding activation for each input position, wherein The layer input comprises an input sequence that is a fragment of an overall sequence derived from the input, and wherein, Each input position is derived from a word in one language, and the network output comprises an output sequence providing a translation of the word into another language, to perform a neural machine translation task, or The input comprises audio data for performing an audio processing task, and the network output provides a result of the audio processing task to recognize words or phrases in the audio, or convert the audio to text, or recognize the natural language in which the spoken utterance is spoken, or The input is text in a natural language or features of text in a natural language, and the network output defines a spectrogram, or other data including audio defining the text spoken in the natural language, to perform a text-to-speech task, or The input is a conditional input and the network output is a sequence of intensity values ​​of pixels of an image to perform image generation, or The input is a sequence of observations or other data characterizing the state of the environment, and the network output defines an action to be performed by the mechanical agent in response to the most recent data in the sequence to perform the agent control task, The method comprises, for each particular network input in the sequence: For each attention layer in the neural network: maintaining corresponding episode memory data, the corresponding episode memory data comprising corresponding hidden states processed by the attention layer for a first portion of previous network inputs preceding the particular network input in the sequence; maintaining corresponding compressed memory data comprising compressed representations of corresponding hidden states processed by the attention layer for a second portion of previous network inputs preceding the first portion of previous network inputs in the sequence; receiving a layer input to be processed by the attention layer during processing of the particular network input using the neural network; and Applying an attention mechanism to (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 corresponding hidden state at each of the plurality of input positions in the particular network input to generate a corresponding activation for each input position in the layer input.

2. The method according to claim 1, further comprising: Updating the episode memory data to include a corresponding hidden state for each input position in the particular network input; as well as The compressed memory data is updated to include a compressed representation of a corresponding hidden state of an earliest network input in the first portion of previous network inputs.

3. The method according to claim 2, wherein: Updating the compressed memory data to include a compressed representation of a corresponding hidden state of an earliest network input in the first portion of previous network inputs includes: Determine the compression ratio; Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input according to the compression rate to generate a compressed representation; and The compressed memory data is modified to include the generated compressed representation.

4. The method according to claim 3, further comprising: The hidden state corresponding to the earliest network input in the first portion of previous network inputs is removed from the episode memory data.

5. The method according to claim 3, wherein Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of the previous network input according to the compression rate includes: Applying a max pooling function to corresponding hidden states of the earliest network input and corresponding hidden states processed by the attention layer for the second portion of previous network input, with a stride equal to the compression rate.

6. The method according to claim 3, wherein: Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of the previous network input according to the compression rate includes: Applying an average pooling function to corresponding hidden states of the earliest network input and corresponding hidden states processed by the attention layer for the second portion of previous network input, with a stride equal to the compression rate.

7. The method according to claim 3, wherein: Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of the previous network input according to the compression rate includes: sorting the corresponding hidden states of the earliest network input and the corresponding hidden states processed by the attention layer for the second portion of previous network input in descending order of corresponding activation weights associated with the plurality of input positions in the earliest network input and the second portion of previous network input; and The corresponding hidden state of the position in the earliest network input and the corresponding hidden state associated with the lowest activation weight processed by the attention layer for the second portion of previous network input are discarded.

8. The method according to claim 3, wherein: Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of the previous network input according to the compression rate includes: Determine the kernel size used for the 1D convolution function; and Applying the 1D convolution function with a stride equal to the compression ratio and a kernel size equal to the determined kernel size to corresponding hidden states of the earliest network input and corresponding hidden states processed by the attention layer for the second portion of previous network input.

9. The method according to claim 3, wherein: Compressing the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of the previous network input according to the compression rate includes: Applying a multi-layer dilated 1D convolution function to the corresponding hidden state of the earliest network input and the corresponding hidden state processed by the attention layer for the second portion of previous network input.

10. The method according to claim 1, wherein The particular network input is an input that is preceded by at least a predetermined threshold number of network inputs in the sequence.

11. The method according to claim 10, further comprising: For each earlier network input of a plurality of earlier network inputs that is preceded by less than the predetermined threshold number of network inputs: For each attention layer in the neural network: maintaining corresponding episode memory data, the corresponding episode memory data comprising respective hidden states to be processed by the attention layer for a previous network input preceding the earlier network input; receiving a layer input to be processed by the attention layer during processing of the earlier network input using the neural network; applying an attention mechanism to (i) a hidden state in the episodic memory data of the layer and (ii) a corresponding hidden state at each of a plurality of input positions in the earlier network input to generate a corresponding activation for each input position in the layer input; and The episode memory data is updated to include a corresponding hidden state for each input position in the earlier network input.

12. The method according to claim 10, wherein: The predetermined threshold number of network inputs is defined by a predetermined threshold value of the size of the episode memory data.

13. The method of claim 2, wherein: The episodic memory data and corresponding hidden state for each input position in the particular network input are represented as corresponding multidimensional arrays; and Updating the episode memory data to include a corresponding hidden state for each input position in the particular network input includes: The two multidimensional arrays are concatenated along a same dimension of the multidimensional arrays.

14. The method according to claim 1, wherein The neural network further includes one or more fully connected layers, one or more layer normalization layers, one or more activation layers or one or more convolutional layers.

15. The method according to any one of claims 1 to 14, further comprising: The neural network is trained on the training data to repeatedly update current values ​​of the network parameters, wherein during training: Gradually increase the number of inputs to train the neural network between each two consecutive updates.

16. The method of claim 15, further comprising, during training: A stopping gradient is applied to (i) the compressed representation in the compressed memory data for the layer and (ii) the hidden state in the episodic memory data for the layer.

17. A system comprising one or more computers and one or more storage devices storing instructions, the instructions being operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the respective methods of any one of claims 1-16.

18. A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the corresponding method of any one of claims 1-16.

Images